Skip to content

Commit

Permalink
Restructure output directory
Browse files Browse the repository at this point in the history
And don't use `atmos.p.output_dir`

The output structure looks like this:
```
coupler_output_dir_amip/
├── checkpoints
│       └── checkpoints for the various models
├── output_0000/
│   ├── atmos/
│   │   └── output of the atmos model
│   └── ocean/
│       └── output of the ocean model
├── output_0001/
│   └── ... component model outputs in their folders ...
├── output_0002/
│   └── ... component model outputs in their folders ...
└── output_active -> output_0002/
```
  • Loading branch information
Sbozzolo committed Dec 3, 2024
1 parent b278a50 commit 3f83442
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 29 deletions.
63 changes: 46 additions & 17 deletions experiments/ClimaEarth/run_amip.jl
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,44 @@ config_dict = merge(parsed_args, config_dict)

comms_ctx = Utilities.get_comms_context(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.
We use `ClimaUtilities.OutputPathGenerator` to manage the output directories.
`OutputPathGenerator` takes care of creating new folders when simulations are
re-run or restarted, so that no data is lost and simulations can be continued.
The output structure looks like this:
```
coupler_output_dir_amip/
├── checkpoints
│ └── checkpoints for the various models
├── output_0000/
│ ├── atmos/
│ │ └── output of the atmos model
│ └── ocean/
│ └── output of the ocean model
├── output_0001/
│ └── ... component model outputs in their folders ...
├── output_0002/
│ └── ... component model outputs in their folders ...
└── output_active -> output_0002/
```
=#

mode_name = config_dict["mode_name"]
COUPLER_OUTPUT_DIR = joinpath(config_dict["coupler_output_dir"], joinpath(mode_name, job_id))
dir_paths = setup_output_dirs(output_dir = COUPLER_OUTPUT_DIR, 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)"

@info(dir_paths.output)

## set unique random seed if desired, otherwise use default
random_seed = config_dict["unique_seed"] ? time_ns() : 1234
Random.seed!(random_seed)
Expand Down Expand Up @@ -157,8 +195,14 @@ end

## get component model dictionaries (if applicable)
atmos_config_dict, config_dict = get_atmos_config_dict(config_dict, job_id)
# Specify atmos output directory to be inside the coupler output directory and disable
# output with counters in atmos (we already have that in ClimaCoupler)
atmos_output_dir = joinpath(dir_paths.output, "clima_atmos")
atmos_config_dict["output_dir"] = atmos_output_dir
atmos_config_dict["output_dir_style"] = "RemovePreexisting"
atmos_config_object = CA.AtmosConfig(atmos_config_dict)


## read in some parsed command line arguments, required by this script
energy_check = config_dict["energy_check"]
const FT = config_dict["FLOAT_TYPE"] == "Float64" ? Float64 : Float32
Expand Down Expand Up @@ -191,19 +235,6 @@ if comms_ctx.device isa ClimaComms.CUDADevice
config_dict["anim"] = false
end

#=
### 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, and `dir_paths.artifacts` is the directory where
the plots (from postprocessing and the conservation checks) of the simulation will be saved. `dir_paths.regrid` is the directory where the regridding
temporary files will be saved.
=#

COUPLER_OUTPUT_DIR = joinpath(config_dict["coupler_output_dir"], joinpath(mode_name, job_id))
dir_paths = setup_output_dirs(output_dir = COUPLER_OUTPUT_DIR, comms_ctx = comms_ctx)
@info "Coupler output directory $(dir_paths.output)"
@info "Coupler artifacts directory $(dir_paths.artifacts)"

@info(dir_paths.output)
config_dict["print_config_dict"] && @info(config_dict)

#=
Expand Down Expand Up @@ -917,7 +948,6 @@ if ClimaComms.iamroot(comms_ctx)

# 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.output

Expand All @@ -939,16 +969,15 @@ if ClimaComms.iamroot(comms_ctx)
# 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")
diagnostics_folder_path = atmos_sim.integrator.p.output_dir
leaderboard_base_path = dir_paths.artifacts
compute_leaderboard(leaderboard_base_path, diagnostics_folder_path)
compute_leaderboard(leaderboard_base_path, atmos_output_dir)
end
end
## plot extra atmosphere diagnostics if specified
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_output_dir, dir_paths.artifacts)
end

## plot all model states and coupler fields (useful for debugging)
Expand Down
41 changes: 29 additions & 12 deletions experiments/ClimaEarth/user_io/io_helpers.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5,35 +5,52 @@
Create output directories for the experiment. If `comms_ctx` is provided, only the root process will create the directories.
By default, the regrid directory is created as a temporary directory inside the output directory,
and the artifacts directory is created inside the output directory with the suffix `_artifacts`.
and the artifacts directory is created inside the output directory with the name `artifacts/`.
`ClimaUtilities.OutputPathGenerator` is used so that simulations can be re-run and re-started.
The output path looks like:
```
coupler_output_dir_amip/
├── checkpoints
│ └── checkpoints for the various models
├── output_0000/
│ ├── atmos/
│ │ └── output of the atmos model
│ └── ocean/
│ └── output of the ocean model
├── output_0001/
│ └── ... component model outputs in their folders ...
├── output_0002/
│ └── ... component model outputs in their folders ...
└── output_active -> output_0002/
```
# Arguments
- `output_dir::String`: The directory where the output files will be stored. Default is the current directory.
- `regrid_dir::String`: The directory where the regridded files will be stored. Default is `output_dir/regrid_tmp/`.
- `checkpoint_dir::String`: The directory where the checkpoint files will be stored. Default is `output_dir/checkpoints/`.
- `artifacts_dir::String`: The directory where the artifacts will be stored. Default is `output_dir_artifacts/`.
- `comms_ctx::Union{Nothing, ClimaComms.AbstractCommsContext}`: The communicator context. If provided, only the root process will create the directories.
# Returns
- A tuple with the paths to the output, regrid, and artifacts directories.
"""
function setup_output_dirs(; output_dir = nothing, artifacts_dir = nothing, comms_ctx)
if output_dir === nothing
output_dir = "."
end
if artifacts_dir === nothing
artifacts_dir = joinpath(output_dir, "artifacts")
end

@info(output_dir)
function setup_output_dirs(;
output_dir = pwd(),
artifacts_dir = joinpath(output_dir, "artifacts"),
checkpoints_dir = joinpath(output_dir, "checkpoints"),
comms_ctx,
)
output_dir = generate_output_path(output_dir)
regrid_dir = nothing
if ClimaComms.iamroot(comms_ctx)
mkpath(output_dir)
mkpath(artifacts_dir)
mkpath(checkpoints_dir)
regrid_dir = mktempdir(output_dir, prefix = "regrid_tmp_")
end
regrid_dir = ClimaComms.bcast(comms_ctx, regrid_dir)

return (; output = output_dir, artifacts = artifacts_dir, regrid = regrid_dir)
return (; output = output_dir, artifacts = artifacts_dir, regrid = regrid_dir, checkpoints = checkpoints_dir)
end

"""
Expand Down

0 comments on commit 3f83442

Please sign in to comment.