use sim_mode only in init

Author: juliasloan25

experiments/ClimaEarth/run_amip.jl

@@ -280,7 +280,6 @@ if sim_mode <: AMIPMode
     ocean_fraction = FT(1) .- ice_fraction .- land_fraction
     ocean_sim = PrescribedOceanSimulation(FT, boundary_space, date0, t_start, ocean_fraction, thermo_params, comms_ctx)
 
-    mode_specifics = (; type = sim_mode)
     Utilities.show_memory_usage()
 
 elseif (sim_mode <: AbstractSlabplanetSimulationMode) && !(sim_mode <: SlabplanetEisenmanMode)
@@ -335,7 +334,6 @@ elseif (sim_mode <: AbstractSlabplanetSimulationMode) && !(sim_mode <: Slabplane
         thermo_params = thermo_params,
     ))
 
-    mode_specifics = (; type = sim_mode)
     Utilities.show_memory_usage()
 
 elseif sim_mode <: SlabplanetEisenmanMode
@@ -382,7 +380,6 @@ elseif sim_mode <: SlabplanetEisenmanMode
         thermo_params = thermo_params,
     )
 
-    mode_specifics = (; type = sim_mode)
     Utilities.show_memory_usage()
 end
 
@@ -487,14 +484,14 @@ end
 #= Set up default AMIP diagnostics
 Use ClimaDiagnostics for default AMIP diagnostics, which currently include turbulent energy fluxes.
 =#
-if sim_mode <: AMIPMode && use_coupler_diagnostics
+if use_coupler_diagnostics
+    @info "Using default coupler diagnostics"
     include("user_io/amip_diagnostics.jl")
     coupler_diags_path = joinpath(dir_paths.output, "coupler")
     isdir(coupler_diags_path) || mkpath(coupler_diags_path)
-    amip_diags_handler =
-        amip_diagnostics_setup(coupler_fields, coupler_diags_path, dates.date0[1], tspan[1], calendar_dt)
+    diags_handler = amip_diagnostics_setup(coupler_fields, coupler_diags_path, dates.date0[1], tspan[1], calendar_dt)
 else
-    amip_diags_handler = nothing
+    diags_handler = nothing
 end
 
 #=
@@ -515,12 +512,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     [tspan[1], tspan[2]],
     Δt_cpl,
     model_sims,
-    mode_specifics,
     callbacks,
     dir_paths,
     turbulent_fluxes,
     thermo_params,
-    amip_diags_handler,
+    diags_handler,
 );
 Utilities.show_memory_usage()
 
@@ -659,8 +655,7 @@ function solve_coupler!(cs)
         ## compute/output AMIP diagnostics if scheduled for this timestep
         ## wrap the current CoupledSimulation fields and time in a NamedTuple to match the ClimaDiagnostics interface
         cs_nt = (; u = cs.fields, p = nothing, t = t, step = round(t / Δt_cpl))
-        (cs.mode.type <: AMIPMode && !isnothing(cs.amip_diags_handler)) &&
-            CD.orchestrate_diagnostics(cs_nt, cs.amip_diags_handler)
+        !isnothing(cs.diags_handler) && CD.orchestrate_diagnostics(cs_nt, cs.diags_handler)
     end
     return nothing
 end
@@ -737,13 +732,7 @@ The postprocessing includes:
 =#
 
 if ClimaComms.iamroot(comms_ctx)
-    postprocessing_vars = (;
-        plot_diagnostics,
-        use_coupler_diagnostics,
-        output_default_diagnostics,
-        t_end,
-        conservation_softfail,
-        atmos_output_dir,
-    )
-    postprocess_sim(cs.mode.type, cs, postprocessing_vars)
+    postprocessing_vars =
+        (; plot_diagnostics, use_coupler_diagnostics, output_default_diagnostics, t_end, conservation_softfail)
+    postprocess_sim(cs, postprocessing_vars)
 end

experiments/ClimaEarth/run_cloudless_aquaplanet.jl

@@ -241,12 +241,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     [tspan[1], tspan[2]],
     Δt_cpl,
     model_sims,
-    (;), # mode_specifics
     callbacks,
     dir_paths,
     turbulent_fluxes,
     thermo_params,
-    nothing, # amip_diags_handler
+    nothing, # diags_handler
 );
 
 #=

experiments/ClimaEarth/run_cloudy_aquaplanet.jl

@@ -265,12 +265,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     [tspan[1], tspan[2]],
     Δt_cpl,
     model_sims,
-    (;), # mode_specifics
     callbacks,
     dir_paths,
     turbulent_fluxes,
     thermo_params,
-    nothing, # amip_diags_handler
+    nothing, # diags_handler
 );
 
 #=

experiments/ClimaEarth/run_cloudy_slabplanet.jl

@@ -301,12 +301,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     [tspan[1], tspan[2]],
     Δt_cpl,
     model_sims,
-    (;), # mode_specifics
     callbacks,
     dir_paths,
     turbulent_fluxes,
     thermo_params,
-    nothing, # amip_diags_handler
+    nothing, # diags_handler
 );
 
 #=

experiments/ClimaEarth/run_dry_held_suarez.jl

@@ -189,12 +189,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     [tspan[1], tspan[2]],
     Δt_cpl,
     model_sims,
-    (;), # mode_specifics
     callbacks,
     dir_paths,
     nothing, # turbulent_fluxes
     thermo_params,
-    nothing, # amip_diags_handler
+    nothing, # diags_handler
 );
 
 #=

experiments/ClimaEarth/run_moist_held_suarez.jl

@@ -239,12 +239,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     [tspan[1], tspan[2]],
     Δt_cpl,
     model_sims,
-    (;), # mode_specifics
     callbacks,
     dir_paths,
     turbulent_fluxes,
     thermo_params,
-    nothing, # amip_diags_handler
+    nothing, # diags_handler
 );
 
 #=

experiments/ClimaEarth/user_io/amip_diagnostics.jl

@@ -48,6 +48,6 @@ function amip_diagnostics_setup(fields, output_dir, start_date, t_start, calenda
 
     # Create the diagnostics handler containing the scheduled diagnostics
     scheduled_diags = [F_turb_energy_diag_sched]
-    amip_diags_handler = CD.DiagnosticsHandler(scheduled_diags, fields, nothing, t_start)
-    return amip_diags_handler
+    diags_handler = CD.DiagnosticsHandler(scheduled_diags, fields, nothing, t_start)
+    return diags_handler
 end

experiments/ClimaEarth/user_io/arg_parsing.jl

@@ -28,7 +28,7 @@ function get_coupler_config()
     job_id = isnothing(job_id) ? string(split(split(config_file, '/')[end], '.')[1]) : job_id
 
     # Read in config dictionary from file, overriding the defaults in `parsed_args`
-    config_dict = merge(parsed_args, YAML.load_file(parsed_args["config_file"]))
+    config_dict = merge(parsed_args, YAML.load_file(config_file))
     config_dict["job_id"] = job_id
     return config_dict
 end

experiments/ClimaEarth/user_io/diagnostics_plots.jl

@@ -147,31 +147,27 @@ end
     make_diagnostics_plots(
         output_path::AbstractString,
         plot_path::AbstractString;
-        short_names::Vector{<:AbstractString} = ["mse", "lr", "edt", "ts"],
-        reduction::String = "average",
         output_prefix = "",
     )
 Create plots for diagnostics. The plots are saved to `plot_path`.
-This is the default plotting function for diagnostics and it can be extended
-to include additional diagnostics.
+This function will plot all variables that have been saved in `output_path`.
 The `reduction` keyword argument should be consistent with the reduction used to save the diagnostics.
 """
-function make_diagnostics_plots(
-    output_path::AbstractString,
-    plot_path::AbstractString;
-    short_names::Vector{<:AbstractString} = ["mse", "lr", "edt", "ts"],
-    reduction::String = "average",
-    output_prefix = "",
-)
-
+function make_diagnostics_plots(output_path::AbstractString, plot_path::AbstractString; output_prefix = "")
     simdir = CAN.SimDir(output_path)
-    if !all(v -> v in CAN.available_vars(simdir), short_names)
-        @warn "Not all variables are available in the output directory $output_path"
-        return
-    end
+    short_names = CAN.available_vars(simdir)
+
+    # Return if there are no variables to plot
+    isempty(short_names) && return
 
     # Create a CAN.OutputVar for each input field
-    vars = [get(simdir; short_name, reduction) for short_name in short_names]
+    vars = Array{CAN.OutputVar}(undef, length(short_names))
+    for (i, short_name) in enumerate(short_names)
+        # Use "average" if available, otherwise use the first reduction
+        reductions = CAN.available_reductions(simdir; short_name)
+        "average" in reductions ? (reduction = "average") : (reduction = first(reductions))
+        vars[i] = get(simdir; short_name, reduction)
+    end
 
     # Filter vars into 2D and 3D variable diagnostics vectors
     # 3D fields are zonally averaged platted on the lat-z plane

experiments/ClimaEarth/user_io/postprocessing.jl

@@ -4,16 +4,40 @@ include("diagnostics_plots.jl")
 include("../leaderboard/leaderboard.jl")
 
 """
-    postprocess_sim(::Type{AbstractSlabplanetSimulationMode}, cs, postprocessing_vars)
+    postprocess_sim(cs, postprocessing_vars)
 
-Call `common_postprocessing` to perform common postprocessing tasks that are common to all simulation types.
-Then, if conservation checks exist, perform them.
+Perform all postprocessing operations. This includes plotting all available
+diagnostics, plotting all model states and coupler fields for debugging,
+producing the leaderboard if monthly data is available, performing
+conservation checks if enabled, and closing all diagnostics file writers.
 """
-function postprocess_sim(::Type{<:AbstractSlabplanetSimulationMode}, cs, postprocessing_vars)
-    (; conservation_softfail,) = postprocessing_vars
+function postprocess_sim(cs, postprocessing_vars)
+    (; use_coupler_diagnostics, output_default_diagnostics, t_end, conservation_softfail) = postprocessing_vars
+    output_dir = cs.dirs.output
+    artifact_dir = cs.dirs.artifacts
+    coupler_output_dir = joinpath(output_dir, "coupler")
+    atmos_output_dir = joinpath(output_dir, "clima_atmos")
+    land_output_dir = joinpath(output_dir, "clima_land")
+
+    # Plot generic diagnostics if requested
+    if use_coupler_diagnostics && plot_diagnostics
+        @info "Plotting diagnostics for coupler, atmos, and land"
+        make_diagnostics_plots(coupler_output_dir, artifact_dir, output_prefix = "coupler_")
+        make_diagnostics_plots(atmos_output_dir, artifact_dir, output_prefix = "atmos_")
+        make_diagnostics_plots(land_output_dir, artifact_dir, output_prefix = "land_")
+    end
+
+    # Plot all model states and coupler fields (useful for debugging)
+    !CA.is_distributed(cs.comms_ctx) && debug(cs, artifact_dir)
 
-    common_postprocessing(cs, postprocessing_vars)
+    # If we have monthly data, plot the leaderboard
+    if t_end > 84600 * 31 * 3 && output_default_diagnostics
+        leaderboard_base_path = artifact_dir
+        compute_leaderboard(leaderboard_base_path, atmos_output_dir)
+        compute_pfull_leaderboard(leaderboard_base_path, atmos_output_dir)
+    end
 
+    # Perform conservation checks if they exist
     if !isnothing(cs.conservation_checks)
         @info "Conservation Check Plots"
         plot_global_conservation(
@@ -31,66 +55,7 @@ function postprocess_sim(::Type{<:AbstractSlabplanetSimulationMode}, cs, postpro
             figname2 = joinpath(cs.dirs.artifacts, "total_water_log_bucket.png"),
         )
     end
-end
-
-"""
-    postprocess_sim(::Type{AMIPMode}, cs, postprocessing_vars)
-
-Call `common_postprocessing` to perform postprocessing tasks that are common to all simulation
-types, and then conditionally plot AMIP diagnostics
-"""
-function postprocess_sim(::Type{AMIPMode}, cs, postprocessing_vars)
-    (; use_coupler_diagnostics, output_default_diagnostics, t_end) = postprocessing_vars
-
-    common_postprocessing(cs, postprocessing_vars)
-
-    if use_coupler_diagnostics
-        ## plot data that correspond to the model's last save_hdf5 call (i.e., last month)
-        @info "AMIP plots"
-        # define variable names and output directories for each diagnostic
-        amip_short_names_atmos = ["ta", "ua", "hus", "clw", "pr", "ts", "toa_fluxes_net"]
-        amip_short_names_coupler = ["F_turb_energy"]
-        output_dir_coupler = cs.dirs.output
-
-        # Check if all output variables are available in the specified directories
-        make_diagnostics_plots(
-            atmos_output_dir,
-            cs.dirs.artifacts,
-            short_names = amip_short_names_atmos,
-            output_prefix = "atmos_",
-        )
-        make_diagnostics_plots(
-            output_dir_coupler,
-            cs.dirs.artifacts,
-            short_names = amip_short_names_coupler,
-            output_prefix = "coupler_",
-        )
-    end
-
-    # Check this because we only want monthly data for making plots
-    if t_end > 84600 * 31 * 3 && output_default_diagnostics
-        leaderboard_base_path = cs.dirs.artifacts
-        compute_leaderboard(leaderboard_base_path, atmos_output_dir)
-        compute_pfull_leaderboard(leaderboard_base_path, atmos_output_dir)
-    end
-
-    # close all AMIP diagnostics file writers
-    !isnothing(cs.amip_diags_handler) &&
-        map(diag -> close(diag.output_writer), cs.amip_diags_handler.scheduled_diagnostics)
-end
-
-"""
-    common_postprocessing(cs, postprocessing_vars)
-
-Perform postprocessing common to all simulation types.
-"""
-function common_postprocessing(cs, postprocessing_vars)
-    (; plot_diagnostics, atmos_output_dir) = postprocessing_vars
-    if plot_diagnostics
-        @info "Plotting diagnostics"
-        make_diagnostics_plots(atmos_output_dir, cs.dirs.artifacts)
-    end
 
-    # plot all model states and coupler fields (useful for debugging)
-    !CA.is_distributed(cs.comms_ctx) && debug(cs, cs.dirs.artifacts)
+    # Close all diagnostics file writers
+    !isnothing(cs.diags_handler) && map(diag -> close(diag.output_writer), cs.diags_handler.scheduled_diagnostics)
 end