Tr/isolate run amip postprocessing (#1120)

* Replace mode_name with mode_type in run_amip This commit adds a struct for each different mode_name. A struct is returned by get_cupler_args as mode_type. This is used everywhere mode_name was. The structs do have hierarchy, but the Abstract types should be given better names. * Isolate run_amp.jl postprocessing This commit adds a new file, experiments/ClimaEarth/user_io/postprocessing.jl, which contains the `postprocess_sim` function. This function copies the if statements previously at the end of run_amip.jl, but it uses dispatch for the simulation mode. * Improve comments and docs in postprocessing.jl * Add docstrings for simulation mode types * Fix NEWS.md formatting * Add News.md section for simulation mode structs and postprocessing * Change simulation mode structs to abstract types This is more in line with the ComponentModelSimulation types in Interfacer.jl * Fix documentation errors
CliMA · Dec 20, 2024 · 958f042 · 958f042
1 parent a29c941
commit 958f042
Show file tree

Hide file tree

Showing 6 changed files with 253 additions and 100 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -42,11 +42,35 @@ pressure levels are being created.
 
 ### Code cleanup
 
-#### Output path updates - PRs [#1058](https://github.com/CliMA/ClimaCoupler.jl/pull/1058),
-    [#1106](https://github.com/CliMA/ClimaCoupler.jl/pull/1106),
-    [#1123](https://github.com/CliMA/ClimaCoupler.jl/pull/1123)
+#### Abstract types representing simulation modes & postprocessing changes - PR [#1120](https://github.com/CliMA/ClimaCoupler.jl/pull/1120)
 
+The available simulation modes are now represented by the following abstract types:
 
+- `ClimaCoupler.Interfacer.AMIPMode`
+- `ClimaCoupler.Interfacer.SlabplanetMode`
+- `ClimaCoupler.Interfacer.SlabplanetAquaMode`
+- `ClimaCoupler.Interfacer.SlabplanetTerraMode`
+- `ClimaCoupler.Interfacer.SlabplanetEisenmanMode`
+
+All of the above types are subtypes of the abstract
+`ClimaCoupler.Interfacer.AbstractSlabplanetSimulationMode`, and all of them except
+`ClimaCoupler.Interfacer.AMIPMode` are subtypes of `ClimaCoupler.Interfacer.AbstractSlabplanetSimulationMode`.
+
+These types are used in `experiments/ClimaEarth/run_amip.jl` instead of representing the
+simulation mode as a string.
+
+The postprocessing in `experiments/ClimaEarth/run_amip.jl` is now moved into functions in
+the new `experiments/ClimaEarth/user_io/postprocessing.jl`. When the simulation is complete,
+`postprocess_sim` is called using the type representing the simulation mode for dispatch.
+`postprocess_sim` has one method for all slabplanet simulation modes and another for the AMIP
+simulation mode. All postprocessing common to all simulation modes is done in the `common_postprocessing`
+function, which is called by both `postprocess_sim` methods.
+
+#### Output path updates - PRs [#1058][1], [#1106][2], [#1123][3]
+
+[1]: https://github.com/CliMA/ClimaCoupler.jl/pull/1058
+[2]: https://github.com/CliMA/ClimaCoupler.jl/pull/1106
+[3]: https://github.com/CliMA/ClimaCoupler.jl/pull/1123
 Previously, ClimaEarth simulation outputs were saved in a path
 `experiments/ClimaEarth/output/$mode_name/$job_id/artifacts/`. Now, `ClimaEarth`
 creates output folders with an increment (increasing the counter every time the
@@ -72,7 +96,7 @@ coupler_output_dir_amip/
 ├── output_0002/
 │   └── ... component model outputs in their folders ...
 └── output_active -> output_0002/
-``
+```
 Note that any external scripts that assume an output path will need to be updated.
 
 #### Remove ClimaCoupler.Diagnostics module - PR [#953](https://github.com/CliMA/ClimaCoupler.jl/pull/953)

diff --git a/docs/src/interfacer.md b/docs/src/interfacer.md
@@ -199,11 +199,23 @@ end
     ClimaCoupler.Interfacer.AtmosModelSimulation
     ClimaCoupler.Interfacer.SurfaceModelSimulation
     ClimaCoupler.Interfacer.ComponentModelSimulation
-    ClimaCoupler.Interfacer.AbstractSimulation
     ClimaCoupler.Interfacer.SurfaceStub
-    ClimaCoupler.Interfacer.stub_init
     ClimaCoupler.Interfacer.float_type
     ClimaCoupler.Interfacer.name
     ClimaCoupler.Interfacer.get_field
     ClimaCoupler.Interfacer.update_field!
+    ClimaCoupler.Interfacer.AbstractSlabplanetSimulationMode
+    ClimaCoupler.Interfacer.AMIPMode
+    ClimaCoupler.Interfacer.SlabplanetMode
+    ClimaCoupler.Interfacer.SlabplanetAquaMode
+    ClimaCoupler.Interfacer.SlabplanetTerraMode
+    ClimaCoupler.Interfacer.SlabplanetEisenmanMode
+```
+
+## Interfacer Internal Functions and Types
+
+```@docs
+    ClimaCoupler.Interfacer.stub_init
+    ClimaCoupler.Interfacer.AbstractSimulation
+    ClimaCoupler.Interfacer.AbstractSimulationMode
 ```
diff --git a/experiments/ClimaEarth/run_amip.jl b/experiments/ClimaEarth/run_amip.jl
@@ -49,6 +49,13 @@ import ClimaCore as CC
 import ClimaCoupler
 import ClimaCoupler:
     ConservationChecker, Checkpointer, FieldExchanger, FluxCalculator, Interfacer, TimeManager, Utilities
+import ClimaCoupler.Interfacer:
+    AbstractSlabplanetSimulationMode,
+    AMIPMode,
+    SlabplanetAquaMode,
+    SlabplanetEisenmanMode,
+    SlabplanetMode,
+    SlabplanetTerraMode
 
 import ClimaUtilities.SpaceVaryingInputs: SpaceVaryingInput
 import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput, evaluate!
@@ -73,9 +80,6 @@ include("components/ocean/slab_ocean.jl")
 include("components/ocean/prescr_seaice.jl")
 include("components/ocean/eisenman_seaice.jl")
 
-## helpers for user-specified IO
-include("user_io/debug_plots.jl")
-
 #=
 ### Configuration Dictionaries
 Each simulation mode has its own configuration dictionary. The `config_dict` of each simulation is a merge of the default configuration
@@ -86,6 +90,7 @@ We can additionally pass the configuration dictionary to the component model ini
 
 include("cli_options.jl")
 include("user_io/arg_parsing.jl")
+include("user_io/postprocessing.jl")
 config_dict = get_coupler_config()
 
 # Select the correct timestep for each component model based on which are available
@@ -95,7 +100,7 @@ add_extra_diagnostics!(config_dict)
 
 (;
     job_id,
-    mode_name,
+    sim_mode,
     random_seed,
     FT,
     comms_ctx,
@@ -244,8 +249,8 @@ In this section of the code, we initialize all component models and read in the
 The specific models and data that are set up depend on which mode we're running.
 =#
 
-@info(mode_name)
-if mode_name == "amip"
+@info(sim_mode)
+if sim_mode <: AMIPMode
     @info("AMIP boundary conditions - do not expect energy conservation")
 
     ## land model
@@ -334,18 +339,18 @@ if mode_name == "amip"
     CO2_field = Interfacer.update_field!(atmos_sim, Val(:co2), CO2_init)
 
     mode_specifics = (;
-        name = mode_name,
+        type = sim_mode,
         SST_timevaryinginput = SST_timevaryinginput,
         SIC_timevaryinginput = SIC_timevaryinginput,
         CO2_timevaryinginput = CO2_timevaryinginput,
     )
     Utilities.show_memory_usage()
 
-elseif mode_name in ("slabplanet", "slabplanet_aqua", "slabplanet_terra")
+elseif (sim_mode <: AbstractSlabplanetSimulationMode) && !(sim_mode <: SlabplanetEisenmanMode)
 
 
-    land_area_fraction = mode_name == "slabplanet_aqua" ? land_area_fraction .* 0 : land_area_fraction
-    land_area_fraction = mode_name == "slabplanet_terra" ? land_area_fraction .* 0 .+ 1 : land_area_fraction
+    land_area_fraction = sim_mode <: SlabplanetAquaMode ? land_area_fraction .* 0 : land_area_fraction
+    land_area_fraction = sim_mode <: SlabplanetTerraMode ? land_area_fraction .* 0 .+ 1 : land_area_fraction
 
     ## land model
     land_sim = bucket_init(
@@ -393,10 +398,10 @@ elseif mode_name in ("slabplanet", "slabplanet_aqua", "slabplanet_terra")
         thermo_params = thermo_params,
     ))
 
-    mode_specifics = (; name = mode_name, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
+    mode_specifics = (; type = sim_mode, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
     Utilities.show_memory_usage()
 
-elseif mode_name == "slabplanet_eisenman"
+elseif sim_mode <: SlabplanetEisenmanMode
 
     ## land model
     land_sim = bucket_init(
@@ -440,7 +445,7 @@ elseif mode_name == "slabplanet_eisenman"
         thermo_params = thermo_params,
     )
 
-    mode_specifics = (; name = mode_name, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
+    mode_specifics = (; type = sim_mode, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
     Utilities.show_memory_usage()
 end
 
@@ -495,7 +500,7 @@ saved in a global `ConservationChecks` struct, `conservation_checks`, which is t
 conservation_checks = nothing
 if energy_check
     @assert(
-        mode_name[1:10] == "slabplanet" && !CA.is_distributed(ClimaComms.context(boundary_space)),
+        sim_mode <: AbstractSlabplanetSimulationMode && !CA.is_distributed(ClimaComms.context(boundary_space)),
         "Only non-distributed slabplanet allowable for energy_check"
     )
     conservation_checks = (;
@@ -537,7 +542,7 @@ albedo_cb = TimeManager.HourlyCallback(
     dt = dt_water_albedo,
     func = FluxCalculator.water_albedo_from_atmosphere!,
     ref_date = [dates.date[1]],
-    active = mode_name == "amip",
+    active = sim_mode <: AMIPMode,
 )
 callbacks =
     (; checkpoint = checkpoint_cb, update_firstdayofmonth! = update_firstdayofmonth!_cb, water_albedo = albedo_cb)
@@ -559,7 +564,7 @@ end
 #= Set up default AMIP diagnostics
 Use ClimaDiagnostics for default AMIP diagnostics, which currently include turbulent energy fluxes.
 =#
-if mode_name == "amip" && use_coupler_diagnostics
+if sim_mode <: AMIPMode && use_coupler_diagnostics
     include("user_io/amip_diagnostics.jl")
     coupler_diags_path = joinpath(dir_paths.output, "coupler")
     isdir(coupler_diags_path) || mkpath(coupler_diags_path)
@@ -688,7 +693,7 @@ function solve_coupler!(cs)
         ## print date on the first of month
         cs.dates.date[1] >= cs.dates.date1[1] && @info(cs.dates.date[1])
 
-        if cs.mode.name == "amip"
+        if cs.mode.type <: AMIPMode
 
             evaluate!(Interfacer.get_field(ocean_sim, Val(:surface_temperature)), cs.mode.SST_timevaryinginput, t)
             evaluate!(Interfacer.get_field(ice_sim, Val(:area_fraction)), cs.mode.SIC_timevaryinginput, t)
@@ -749,7 +754,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.name == "amip" && !isnothing(cs.amip_diags_handler)) &&
+        (cs.mode.type <: AMIPMode && !isnothing(cs.amip_diags_handler)) &&
             CD.orchestrate_diagnostics(cs_nt, cs.amip_diags_handler)
     end
     return nothing
@@ -832,78 +837,13 @@ The postprocessing includes:
 =#
 
 if ClimaComms.iamroot(comms_ctx)
-
-    ## energy check plots
-    if !isnothing(cs.conservation_checks) && cs.mode.name[1:10] == "slabplanet"
-        @info "Conservation Check Plots"
-        plot_global_conservation(
-            cs.conservation_checks.energy,
-            cs,
-            conservation_softfail,
-            figname1 = joinpath(dir_paths.artifacts, "total_energy_bucket.png"),
-            figname2 = joinpath(dir_paths.artifacts, "total_energy_log_bucket.png"),
-        )
-        plot_global_conservation(
-            cs.conservation_checks.water,
-            cs,
-            conservation_softfail,
-            figname1 = joinpath(dir_paths.artifacts, "total_water_bucket.png"),
-            figname2 = joinpath(dir_paths.artifacts, "total_water_log_bucket.png"),
-        )
-    end
-
-    ## plotting AMIP results
-    if cs.mode.name == "amip"
-        if use_coupler_diagnostics
-            ## plot data that correspond to the model's last save_hdf5 call (i.e., last month)
-            @info "AMIP plots"
-
-            ## ClimaESM
-            include("user_io/diagnostics_plots.jl")
-
-            # 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 = dir_paths.output
-
-            # Check if all output variables are available in the specified directories
-            make_diagnostics_plots(
-                atmos_output_dir,
-                dir_paths.artifacts,
-                short_names = amip_short_names_atmos,
-                output_prefix = "atmos_",
-            )
-            make_diagnostics_plots(
-                output_dir_coupler,
-                dir_paths.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
-            include("leaderboard/leaderboard.jl")
-            leaderboard_base_path = dir_paths.artifacts
-            compute_leaderboard(leaderboard_base_path, atmos_output_dir)
-            compute_pfull_leaderboard(leaderboard_base_path, atmos_output_dir)
-        end
-    end
-    ## plot extra atmosphere diagnostics if specified
-    if plot_diagnostics
-        @info "Plotting diagnostics"
-        include("user_io/diagnostics_plots.jl")
-        make_diagnostics_plots(atmos_output_dir, dir_paths.artifacts)
-    end
-
-    ## plot all model states and coupler fields (useful for debugging)
-    !CA.is_distributed(comms_ctx) && debug(cs, dir_paths.artifacts)
-
-    # if isinteractive() #hide
-    #     ## clean up for interactive runs, retain all output otherwise #hide
-    #     rm(dir_paths.output; recursive = true, force = true) #hide
-    # end #hide
-
-    ## close all AMIP diagnostics file writers
-    !isnothing(amip_diags_handler) && map(diag -> close(diag.output_writer), amip_diags_handler.scheduled_diagnostics)
+    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)
 end
diff --git a/experiments/ClimaEarth/user_io/arg_parsing.jl b/experiments/ClimaEarth/user_io/arg_parsing.jl
@@ -1,5 +1,13 @@
 import YAML
 
+mode_name_dict = Dict(
+    "amip" => AMIPMode,
+    "slabplanet" => SlabplanetMode,
+    "slabplanet_aqua" => SlabplanetAquaMode,
+    "slabplanet_terra" => SlabplanetTerraMode,
+    "slabplanet_eisenman" => SlabplanetEisenmanMode,
+)
+
 """
     get_coupler_config()
 
@@ -42,6 +50,7 @@ function get_coupler_args(config_dict::Dict)
     config_dict["print_config_dict"] && @info(config_dict)
     job_id = config_dict["job_id"]
     mode_name = config_dict["mode_name"]
+    sim_mode = mode_name_dict[mode_name]
 
     # Computational simulation setup information
     random_seed = config_dict["unique_seed"] ? time_ns() : 1234
@@ -91,7 +100,7 @@ function get_coupler_args(config_dict::Dict)
 
     return (;
         job_id,
-        mode_name,
+        sim_mode,
         random_seed,
         FT,
         comms_ctx,