minor dir_paths cleanup

Author: juliasloan25

experiments/ClimaEarth/setup_run.jl

@@ -143,40 +143,29 @@ function CoupledSimulation(config_dict::AbstractDict)
     ) = get_coupler_args(config_dict)
 
     #=
-    ### I/O Directory Setup `setup_output_dirs` returns `dir_paths.output =
-    COUPLER_OUTPUT_DIR`, which is the directory where the output of the simulation
-    will be saved, `dir_paths.artifacts` is the directory where the plots (from
-    postprocessing and the conservation checks) of the simulation will be saved,
-    #and `dir_paths.checkpoints`, where restart files are saved.
+    ### I/O Directory Setup
+    `setup_output_dirs` returns a NamedTuple with the paths to
+    output directories for each component and the coupler,
+    as well as paths to artifacts, regrid, and checkpoints directories.
     =#
-
-    dir_paths =
-        Utilities.setup_output_dirs(output_dir = output_dir_root, comms_ctx = comms_ctx)
-    @info "Coupler output directory $(dir_paths.output)"
-    @info "Coupler artifacts directory $(dir_paths.artifacts)"
-    @info "Coupler checkpoint directory $(dir_paths.checkpoints)"
-
-    atmos_output_dir = joinpath(dir_paths.output, "clima_atmos")
-    isdir(atmos_output_dir) || mkpath(atmos_output_dir)
-    land_output_dir = joinpath(dir_paths.output, "clima_land")
-    isdir(land_output_dir) || mkpath(land_output_dir)
-    ocean_output_dir = joinpath(dir_paths.output, "clima_ocean")
-    isdir(ocean_output_dir) || mkpath(ocean_output_dir)
-
+    dir_paths = Utilities.setup_output_dirs(
+        output_dir_root = output_dir_root,
+        comms_ctx = comms_ctx,
+    )
 
     ## get component model dictionaries (if applicable)
     ## Note this step must come after parsing the coupler config dictionary, since
     ##  some parameters are passed from the coupler config to the component model configs
-    atmos_config_dict = get_atmos_config_dict(config_dict, atmos_output_dir)
+    atmos_config_dict = get_atmos_config_dict(config_dict, dir_paths.atmos_output_dir)
 
     ## set unique random seed if desired, otherwise use default
     Random.seed!(random_seed)
     @info "Random seed set to $(random_seed)"
 
     if detect_restart_files
         isnothing(restart_t) &&
-            (restart_t = Checkpointer.t_start_from_checkpoint(dir_paths.checkpoints))
-        isnothing(restart_dir) && (restart_dir = dir_paths.checkpoints)
+            (restart_t = Checkpointer.t_start_from_checkpoint(dir_paths.checkpoints_dir))
+        isnothing(restart_dir) && (restart_dir = dir_paths.checkpoints_dir)
     end
     should_restart = !isnothing(restart_t) && !isnothing(restart_dir)
     if should_restart
@@ -205,12 +194,6 @@ function CoupledSimulation(config_dict::AbstractDict)
 
     tspan = (t_start, t_end)
 
-    #=
-    ## Data File Paths
-    =#
-    land_mask_data =
-        joinpath(@clima_artifact("landsea_mask_60arcseconds", comms_ctx), "landsea_mask.nc")
-
     #=
     ## Component Model Initialization
     Here we set initial and boundary conditions for each component model. Each component model is required to have an `init` function that
@@ -275,6 +258,8 @@ function CoupledSimulation(config_dict::AbstractDict)
     =#
 
     # Preprocess the file to be 1s and 0s before remapping into onto the grid
+    land_mask_data =
+        joinpath(@clima_artifact("landsea_mask_60arcseconds", comms_ctx), "landsea_mask.nc")
     land_fraction = SpaceVaryingInput(land_mask_data, "landsea", boundary_space)
     land_fraction = ifelse.(land_fraction .> eps(FT), FT(1), FT(0))
 
@@ -325,7 +310,7 @@ function CoupledSimulation(config_dict::AbstractDict)
                 dt = component_dt_dict["dt_land"],
                 tspan,
                 start_date,
-                output_dir = land_output_dir,
+                output_dir = dir_paths.land_output_dir,
                 area_fraction = land_fraction,
                 shared_surface_space,
                 surface_elevation,
@@ -342,7 +327,7 @@ function CoupledSimulation(config_dict::AbstractDict)
                 dt = component_dt_dict["dt_land"],
                 tspan,
                 start_date,
-                output_dir = land_output_dir,
+                output_dir = dir_paths.land_output_dir,
                 area_fraction = land_fraction,
                 shared_surface_space,
                 land_spun_up_ic,
@@ -382,7 +367,7 @@ function CoupledSimulation(config_dict::AbstractDict)
                 ocean_fraction,
                 start_date,
                 stop_date;
-                output_dir = ocean_output_dir,
+                output_dir = dir_paths.ocean_output_dir,
                 comms_ctx,
             )
         else
@@ -410,7 +395,7 @@ function CoupledSimulation(config_dict::AbstractDict)
             dt = component_dt_dict["dt_land"],
             tspan,
             start_date,
-            output_dir = land_output_dir,
+            output_dir = dir_paths.land_output_dir,
             area_fraction = land_fraction,
             surface_elevation,
             land_temperature_anomaly,
@@ -450,9 +435,8 @@ function CoupledSimulation(config_dict::AbstractDict)
 
     ## coupler exchange fields
     coupler_field_names = Interfacer.default_coupler_fields()
-    for sim in model_sims
-        Interfacer.add_coupler_fields!(coupler_field_names, sim)
-    end
+    foreach(sim -> Interfacer.add_coupler_fields!(coupler_field_names, sim), model_sims)
+
     # add coupler fields required to track conservation, if specified
     energy_check && push!(coupler_field_names, :P_net)
 
@@ -500,11 +484,9 @@ function CoupledSimulation(config_dict::AbstractDict)
     =#
     if use_coupler_diagnostics
         @info "Using default coupler diagnostics"
-        coupler_diags_path = joinpath(dir_paths.output, "coupler")
-        isdir(coupler_diags_path) || mkpath(coupler_diags_path)
         diags_handler = coupler_diagnostics_setup(
             coupler_fields,
-            coupler_diags_path,
+            dir_paths.coupler_output_dir,
             start_date,
             tspan[1],
             diagnostics_dt,
@@ -631,28 +613,23 @@ end
 
 Process the results after a simulation has completed, including generating
 plots, checking conservation, and other diagnostics.
+All postprocessing is performed using the root process only, if applicable.
 
 When `conservation_softfail` is true, throw an error if conservation is not
 respected.
 
 When `rmse_check` is true, compute the RMSE against observations and test
 that it is below a certain threshold.
+
+The postprocessing includes:
+- Energy and water conservation checks (if running SlabPlanet with checks enabled)
+- Animations (if not running in MPI)
+- AMIP plots of the final state of the model
+- Error against observations
+- Optional additional atmosphere diagnostics plots
+- Plots of useful coupler and component model fields for debugging
 """
 function postprocess(cs; conservation_softfail = false, rmse_check = false)
-    #=
-    ## Postprocessing
-    All postprocessing is performed using the root process only, if applicable.
-    Our postprocessing consists of outputting a number of plots to visualize the model output.
-
-    The postprocessing includes:
-    - Energy and water conservation checks (if running SlabPlanet with checks enabled)
-    - Animations (if not running in MPI)
-    - AMIP plots of the final state of the model
-    - Error against observations
-    - Optional additional atmosphere diagnostics plots
-    - Plots of useful coupler and component model fields for debugging
-    =#
-
     if ClimaComms.iamroot(ClimaComms.context(cs)) && !isnothing(cs.diags_handler)
         postprocessing_vars = (; conservation_softfail, rmse_check)
         postprocess_sim(cs, postprocessing_vars)

experiments/ClimaEarth/test/component_model_tests/climaland_tests.jl

@@ -97,7 +97,7 @@ end
         Ref(-1), # prev_checkpoint_t
         model_sims,
         (;), # callbacks
-        (;), # dirs
+        (;), # dir_paths
         thermo_params, # thermo_params
         nothing, # diags_handler
     )

experiments/ClimaEarth/test/debug_plots_tests.jl

@@ -87,7 +87,7 @@ plot_field_names(sim::Interfacer.SurfaceStub) = (:stub_field,)
         Ref(-1), # prev_checkpoint_t
         model_sims, # model_sims
         (;), # callbacks
-        (;), # dirs
+        (;), # dir_paths
         nothing, # thermo_params
         nothing, # diags_handler
     )

experiments/ClimaEarth/test/restart.jl

@@ -59,7 +59,7 @@ four_steps_reading = deepcopy(four_steps)
 
 four_steps_reading["t_end"] = "900secs"
 four_steps_reading["detect_restart_files"] = true
-four_steps_reading["restart_dir"] = cs_four_steps.dirs.checkpoints
+four_steps_reading["restart_dir"] = cs_four_steps.dir_paths.checkpoints_dir
 four_steps_reading["restart_t"] = 720
 four_steps_reading["job_id"] = "four_steps_reading"
 

experiments/ClimaEarth/user_io/postprocessing.jl

@@ -14,22 +14,23 @@ conservation checks if enabled, and closing all diagnostics file writers.
 """
 function postprocess_sim(cs, postprocessing_vars)
     (; conservation_softfail, rmse_check) = 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")
-    ocean_output_dir = joinpath(output_dir, "clima_ocean")
+    (;
+        coupler_output_dir,
+        atmos_output_dir,
+        land_output_dir,
+        ocean_output_dir,
+        artifacts_dir,
+    ) = cs.dir_paths
 
     # Plot generic diagnostics
     @info "Plotting diagnostics for coupler, atmos, land, and ocean"
-    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_")
-    make_ocean_diagnostics_plots(ocean_output_dir, artifact_dir, output_prefix = "ocean_")
+    make_diagnostics_plots(coupler_output_dir, artifacts_dir, output_prefix = "coupler_")
+    make_diagnostics_plots(atmos_output_dir, artifacts_dir, output_prefix = "atmos_")
+    make_diagnostics_plots(land_output_dir, artifacts_dir, output_prefix = "land_")
+    make_ocean_diagnostics_plots(ocean_output_dir, artifacts_dir, output_prefix = "ocean_")
 
     # Plot all model states and coupler fields (useful for debugging)
-    ClimaComms.context(cs) isa ClimaComms.SingletonCommsContext && debug(cs, artifact_dir)
+    ClimaComms.context(cs) isa ClimaComms.SingletonCommsContext && debug(cs, artifacts_dir)
 
     # If we have enough data (in time, but also enough variables), plot the leaderboard.
     # We need pressure to compute the leaderboard.
@@ -39,7 +40,7 @@ function postprocess_sim(cs, postprocessing_vars)
         times = CAN.times(get(simdir, first(CAN.available_vars(simdir))))
         t_end = times[end]
         if t_end > 84600 * 31 * 3 # 3 months for spin up
-            leaderboard_base_path = artifact_dir
+            leaderboard_base_path = artifacts_dir
             compute_leaderboard(leaderboard_base_path, atmos_output_dir, 3)
             rmse_check && test_rmse_thresholds(atmos_output_dir, 3)
             pressure_in_output &&
@@ -54,15 +55,15 @@ function postprocess_sim(cs, postprocessing_vars)
             cs.conservation_checks.energy,
             cs,
             conservation_softfail,
-            figname1 = joinpath(cs.dirs.artifacts, "total_energy_bucket.png"),
-            figname2 = joinpath(cs.dirs.artifacts, "total_energy_log_bucket.png"),
+            figname1 = joinpath(artifacts_dir, "total_energy_bucket.png"),
+            figname2 = joinpath(artifacts_dir, "total_energy_log_bucket.png"),
         )
         plot_global_conservation(
             cs.conservation_checks.water,
             cs,
             conservation_softfail,
-            figname1 = joinpath(cs.dirs.artifacts, "total_water_bucket.png"),
-            figname2 = joinpath(cs.dirs.artifacts, "total_water_log_bucket.png"),
+            figname1 = joinpath(artifacts_dir, "total_water_bucket.png"),
+            figname2 = joinpath(artifacts_dir, "total_water_log_bucket.png"),
         )
     end
 
@@ -97,25 +98,29 @@ end
 """
     save_sypd_walltime_to_disk(cs, walltime)
 
-Save the computed `sypd`, `walltime_per_coupling_step`, and memory usage to text files.
+Save the computed `sypd`, `walltime_per_coupling_step`,
+and memory usage to text files in the `artifacts` directory.
 """
 function save_sypd_walltime_to_disk(cs, walltime)
     if ClimaComms.iamroot(ClimaComms.context(cs))
         sypd = simulated_years_per_day(cs, walltime)
         walltime_per_step = walltime_per_coupling_step(cs, walltime)
 
-        open(joinpath(cs.dirs.artifacts, "sypd.txt"), "w") do sypd_filename
+        open(joinpath(cs.dir_paths.artifacts_dir, "sypd.txt"), "w") do sypd_filename
             println(sypd_filename, "$sypd")
         end
 
         open(
-            joinpath(cs.dirs.artifacts, "walltime_per_step.txt"),
+            joinpath(cs.dir_paths.artifacts_dir, "walltime_per_step.txt"),
             "w",
         ) do walltime_per_step_filename
             println(walltime_per_step_filename, "$(walltime_per_step)")
         end
 
-        open(joinpath(cs.dirs.artifacts, "max_rss_cpu.txt"), "w") do cpu_max_rss_filename
+        open(
+            joinpath(cs.dir_paths.artifacts_dir, "max_rss_cpu.txt"),
+            "w",
+        ) do cpu_max_rss_filename
             cpu_max_rss_GB = Utilities.show_memory_usage()
             println(cpu_max_rss_filename, cpu_max_rss_GB)
         end

src/Checkpointer.jl

@@ -134,7 +134,7 @@ function checkpoint_sims(cs::Interfacer.CoupledSimulation)
     time = Int(round(float(cs.t[])))
     day = floor(Int, time / (60 * 60 * 24))
     sec = floor(Int, time % (60 * 60 * 24))
-    output_dir = cs.dirs.checkpoints
+    output_dir = cs.dir_paths.checkpoints_dir
     prev_checkpoint_t = cs.prev_checkpoint_t[]
     comms_ctx = ClimaComms.context(cs)
 

src/Interfacer.jl

@@ -70,7 +70,7 @@ struct CoupledSimulation{
     prev_checkpoint_t::CTT
     model_sims::NTMS
     callbacks::CALLBACKS
-    dirs::NTP
+    dir_paths::NTP
     thermo_params::TP
     diags_handler::DH
 end
@@ -83,7 +83,7 @@ function Base.show(io::IO, sim::CoupledSimulation)
         io,
         "Coupled Simulation\n",
         "├── Running on: $(device_type)\n",
-        "├── Output folder: $(sim.dirs.output)\n",
+        "├── Output folder: $(sim.dir_paths.output_dir_root)\n",
         "└── Current date: $(current_date(sim))\n",
     )
 end