Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Tr/isolate run amip postprocessing #1120

Merged
merged 8 commits into from
Dec 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 28 additions & 4 deletions NEWS.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,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
Comment on lines +66 to +70
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's nice to keep these in the [...] (...) format because it creates a hyperlink (e.g. for PR 1058 here). However as it is in current main, the 2nd and 3rd PRs don't get linked because they're on newlines. Maybe we could do something like this:

Suggested change
#### 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
#### 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)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The notation I used also creates hyperlinks. See the preview here

I prefer this notation because it avoids making a very long line that wraps around when the source is viewed on most screens. If you prefer the normal hyperlink notation, I'll make the change.

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
Expand All @@ -69,7 +93,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)
Expand Down
16 changes: 14 additions & 2 deletions docs/src/interfacer.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
```
124 changes: 32 additions & 92 deletions experiments/ClimaEarth/run_amip.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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!
Expand All @@ -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
Expand All @@ -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
Expand All @@ -95,7 +100,7 @@ add_extra_diagnostics!(config_dict)

(;
job_id,
mode_name,
sim_mode,
random_seed,
FT,
comms_ctx,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This field isn't used anymore (all uses of cs.mode.name are removed), so we can remove it from mode_specifics

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the field is used inside of solver_coupler lines 691 and 752
cs.mode.type isa AMIPMode

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh yeah, good point

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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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 = (;
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
11 changes: 10 additions & 1 deletion experiments/ClimaEarth/user_io/arg_parsing.jl
Original file line number Diff line number Diff line change
@@ -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()

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -91,7 +100,7 @@ function get_coupler_args(config_dict::Dict)

return (;
job_id,
mode_name,
sim_mode,
random_seed,
FT,
comms_ctx,
Expand Down
Loading
Loading