diff --git a/REFACTORING_PR_BREAKDOWN.md b/REFACTORING_PR_BREAKDOWN.md new file mode 100644 index 0000000000..c2d3f9042d --- /dev/null +++ b/REFACTORING_PR_BREAKDOWN.md @@ -0,0 +1,70 @@ +# Refactoring PR Breakdown + +This document outlines how to break down the large refactoring into smaller, manageable PRs. + +## Suggested PR Breakdown: + +### PR 1: Move `restore.jl` to Checkpointer +- Move `restore!` functions to `Checkpointer.jl` +- Update includes in `climaland_bucket.jl` and `climaatmos.jl` +- Update calls to `Checkpointer.restore!` +- **Why first**: Self-contained, minimal dependencies + +### PR 2: Create Input module and move `cli_options.jl` +- Create `Input.jl` module +- Move `argparse_settings()` and `parse_commandline()` from `cli_options.jl` +- Update references to use `Input.argparse_settings` and `Input.parse_commandline` +- **Note**: Can keep `cli_options.jl` temporarily for backward compatibility (or delete if nothing else uses it) + +### PR 3: Move `arg_parsing.jl` functions to Input +- Move `get_coupler_config_dict()` and `get_coupler_args()` to `Input.jl` +- Update `setup_run.jl` to use `Input.get_coupler_config_dict` and `Input.get_coupler_args` +- Update other files that use these functions +- Delete `arg_parsing.jl` after migration + +### PR 4: Create Postprocessor module structure +- Create `Postprocessor.jl` with basic structure +- Move `postprocess()` and `postprocess_sim()` from `setup_run.jl` +- Move `simulated_years_per_day`, `walltime_per_coupling_step`, `save_sypd_walltime_to_disk` +- Update `SimCoordinator` to call `Postprocessor.simulated_years_per_day`, etc. +- **Why now**: Establishes the module without the large files + +### PR 5: Move diagnostics functions to Postprocessor +- Move `coupler_diagnostics.jl` functions to `Postprocessor.jl` +- Update `SimCoordinator` to use `Postprocessor.CD.orchestrate_diagnostics` +- **Why now**: Self-contained, relatively small + +### PR 6: Move plotting functions to Postprocessor +- Create `Postprocessor/` directory +- Move `diagnostics_plots.jl` and `debug_plots.jl` to `Postprocessor/` +- Update includes in `Postprocessor.jl` +- Update `postprocess_sim` to call these functions + +### PR 7: Move leaderboard functions to Postprocessor +- Move `leaderboard/` directory to `Postprocessor/leaderboard/` +- Update includes in `Postprocessor.jl` +- Update `postprocess_sim` to call leaderboard functions + +### PR 8: Move benchmarks to Postprocessor +- Move `benchmarks.jl` to `Postprocessor/` +- Update includes in `Postprocessor.jl` + +### PR 9: Move Postprocessor.jl into Postprocessor/ folder +- Move `Postprocessor.jl` to `Postprocessor/Postprocessor.jl` +- Update include path in `ClimaCoupler.jl` +- Cleanup + +### PR 10: Cleanup and final touches +- Delete old files (`cli_options.jl`, etc.) +- Update any remaining references +- Move `get_field` methods to appropriate component files +- Final cleanup + +## Tips for Managing Interdependencies: + +1. **Use feature flags or temporary compatibility shims** if needed +2. **Keep old files temporarily** and add deprecation warnings +3. **Test each PR independently** - each should be functional on its own +4. **Update imports gradually** - you can keep both old and new imports working temporarily + +The key is to move in dependency order: start with leaf modules (Checkpointer), then modules that depend on them (Input), then modules that depend on those (Postprocessor). diff --git a/experiments/ClimaEarth/components/atmosphere/climaatmos.jl b/experiments/ClimaEarth/components/atmosphere/climaatmos.jl index 557c820832..5a02955428 100644 --- a/experiments/ClimaEarth/components/atmosphere/climaatmos.jl +++ b/experiments/ClimaEarth/components/atmosphere/climaatmos.jl @@ -20,7 +20,6 @@ if pkgversion(CA) < v"0.28.6" CC.Adapt.@adapt_structure CA.RRTMGPInterface.RRTMGPModel end -include("../shared/restore.jl") ### ### Functions required by ClimaCoupler.jl for an AtmosModelSimulation @@ -140,7 +139,7 @@ end function Checkpointer.restore_cache!(sim::ClimaAtmosSimulation, new_cache) comms_ctx = ClimaComms.context(sim.integrator.u.c) - restore!( + Checkpointer.restore!( Checkpointer.get_model_cache(sim), new_cache, comms_ctx; @@ -735,3 +734,23 @@ function climaatmos_restart_path(output_dir_root, t) end error("Restart file for time $t not found") end + +### +### Additional accessor functions for debugging ClimaAtmosSimulation +### + +# Helper function for specific humidity +specific_humidity(::CA.DryModel, integrator) = [eltype(integrator.u)(0)] +specific_humidity(::Union{CA.EquilMoistModel, CA.NonEquilMoistModel}, integrator) = + integrator.u.c.ρq_tot + +# Additional debug fields for ClimaAtmosSimulation +Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:ρq_tot}) = + specific_humidity(sim.integrator.p.atmos.moisture_model, sim.integrator) +Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:ρe_tot}) = sim.integrator.u.c.ρe_tot + +# Plot field names for ClimaAtmosSimulation +# TODO is this the right name? where is this used? +function plot_field_names(sim::ClimaAtmosSimulation) + return (:w, :ρq_tot, :ρe_tot, :liquid_precipitation, :snow_precipitation) +end diff --git a/experiments/ClimaEarth/components/land/climaland_bucket.jl b/experiments/ClimaEarth/components/land/climaland_bucket.jl index 5587d662db..410f324c1a 100644 --- a/experiments/ClimaEarth/components/land/climaland_bucket.jl +++ b/experiments/ClimaEarth/components/land/climaland_bucket.jl @@ -13,7 +13,6 @@ import ClimaCoupler: Checkpointer, FluxCalculator, Interfacer, FieldExchanger using NCDatasets include("climaland_helpers.jl") -include("../shared/restore.jl") ### ### Functions required by ClimaCoupler.jl for a SurfaceModelSimulation @@ -411,7 +410,7 @@ end function Checkpointer.restore_cache!(sim::BucketSimulation, new_cache) old_cache = Checkpointer.get_model_cache(sim) comms_ctx = ClimaComms.context(sim.model) - restore!( + Checkpointer.restore!( old_cache, new_cache, comms_ctx, diff --git a/experiments/ClimaEarth/components/land/climaland_integrated.jl b/experiments/ClimaEarth/components/land/climaland_integrated.jl index 7e60ad0b0e..6d62219fa7 100644 --- a/experiments/ClimaEarth/components/land/climaland_integrated.jl +++ b/experiments/ClimaEarth/components/land/climaland_integrated.jl @@ -533,7 +533,7 @@ end function Checkpointer.restore_cache!(sim::ClimaLandSimulation, new_cache) old_cache = Checkpointer.get_model_cache(sim) comms_ctx = ClimaComms.context(sim.model.soil) - restore!( + Checkpointer.restore!( old_cache, new_cache, comms_ctx, diff --git a/experiments/ClimaEarth/input/bucket_ic_august.nc b/experiments/ClimaEarth/input/bucket_ic_august.nc deleted file mode 100644 index 2f4295c934..0000000000 Binary files a/experiments/ClimaEarth/input/bucket_ic_august.nc and /dev/null differ diff --git a/experiments/ClimaEarth/setup_run.jl b/experiments/ClimaEarth/setup_run.jl index a670cbc79a..73f4a2fe25 100644 --- a/experiments/ClimaEarth/setup_run.jl +++ b/experiments/ClimaEarth/setup_run.jl @@ -37,7 +37,9 @@ import ClimaCoupler: Checkpointer, FieldExchanger, FluxCalculator, + Input, Interfacer, + Postprocessor, TimeManager, Utilities import ClimaCoupler.Interfacer: @@ -88,10 +90,6 @@ dictionary and the simulation-specific configuration dictionary, which allows th We can additionally pass the configuration dictionary to the component model initializers, which will then override the default settings of the component models. =# -include("cli_options.jl") -include("user_io/arg_parsing.jl") -include("user_io/postprocessing.jl") -include("user_io/coupler_diagnostics.jl") """ CoupledSimulation(config_file) @@ -106,7 +104,7 @@ needed to run a coupled simulation. function CoupledSimulation( config_file = joinpath(pkgdir(ClimaCoupler), "config/ci_configs/amip_default.yml"), ) - config_dict = get_coupler_config_dict(config_file) + config_dict = Input.get_coupler_config_dict(config_file) return CoupledSimulation(config_dict) end @@ -146,7 +144,7 @@ function CoupledSimulation(config_dict::AbstractDict) parameter_files, era5_initial_condition_dir, ice_model, - ) = get_coupler_args(config_dict) + ) = Input.get_coupler_args(config_dict) # Get default shared parameters from ClimaParams.jl, overriding with any provided parameter files override_file = CP.merge_toml_files(parameter_files; override = true) @@ -596,89 +594,6 @@ function CoupledSimulation(config_dict::AbstractDict) return cs end -""" - run!(cs::CoupledSimulation) - -Evolve the given simulation, producing plots and other diagnostic information. - -Keyword arguments -================== - -`precompile`: If `true`, run the coupled simulations for two steps, so that most functions - are precompiled and subsequent timing will be more accurate. -""" -function run!( - cs::CoupledSimulation; - precompile = (cs.tspan[end] > 2 * cs.Δt_cpl + cs.tspan[begin]), -) - - ## Precompilation of Coupling Loop - # Here we run the entire coupled simulation for two timesteps to precompile several - # functions for more accurate timing of the overall simulation. - precompile && (step!(cs); step!(cs)) - - ## Run garbage collection before solving for more accurate memory comparison to ClimaAtmos - GC.gc() - - #= - ## Solving and Timing the Full Simulation - - This is where the full coupling loop, `solve_coupler!` is called for the full timespan of the simulation. - We use the `ClimaComms.@elapsed` macro to time the simulation on both CPU and GPU, and use this - value to calculate the simulated years per day (SYPD) of the simulation. - =# - @info "Starting coupling loop" - walltime = ClimaComms.@elapsed ClimaComms.device(cs) begin - s = CA.@timed_str begin - while cs.t[] < cs.tspan[end] - step!(cs) - end - end - end - @info "Simulation took $(walltime) seconds" - - sypd = simulated_years_per_day(cs, walltime) - walltime_per_step = walltime_per_coupling_step(cs, walltime) - @info "SYPD: $sypd" - @info "Walltime per coupling step: $(walltime_per_step)" - save_sypd_walltime_to_disk(cs, walltime) - - # Close all diagnostics file writers - isnothing(cs.diags_handler) || - foreach(diag -> close(diag.output_writer), cs.diags_handler.scheduled_diagnostics) - foreach(Interfacer.close_output_writers, cs.model_sims) - - return nothing -end - -""" - postprocess(cs; conservation_softfail = false, rmse_check = false) - -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) - if ClimaComms.iamroot(ClimaComms.context(cs)) && !isnothing(cs.diags_handler) - postprocessing_vars = (; conservation_softfail, rmse_check) - postprocess_sim(cs, postprocessing_vars) - end - return nothing -end """ setup_and_run(config_dict) @@ -701,43 +616,3 @@ function setup_and_run(config_dict) run!(cs) return cs end - -""" - step!(cs::CoupledSimulation) - -Take one coupling step forward in time. - -This function runs the component models sequentially, and exchanges combined fields and -calculates fluxes using the selected turbulent fluxes option. Note, one coupling step might -require multiple steps in some of the component models. -""" -function step!(cs::CoupledSimulation) - # Update the current time - cs.t[] += cs.Δt_cpl - - ## compute global energy and water conservation checks - ## (only for slabplanet if tracking conservation is enabled) - ConservationChecker.check_conservation!(cs) - - ## step component model simulations sequentially for one coupling timestep (Δt_cpl) - FieldExchanger.step_model_sims!(cs) - - ## update the surface fractions for surface models - FieldExchanger.update_surface_fractions!(cs) - - ## exchange all non-turbulent flux fields between models, including radiative and precipitation fluxes - FieldExchanger.exchange!(cs) - - ## calculate turbulent fluxes in the coupler and update the model simulations with them - FluxCalculator.turbulent_fluxes!(cs) - - ## compute any ocean-sea ice fluxes - FluxCalculator.ocean_seaice_fluxes!(cs) - - ## Maybe call the callbacks - TimeManager.callbacks!(cs) - - # Compute coupler diagnostics - CD.orchestrate_diagnostics(cs) - return nothing -end diff --git a/experiments/ClimaEarth/user_io/postprocessing.jl b/experiments/ClimaEarth/user_io/postprocessing.jl index 54b043deac..ea1b94630d 100644 --- a/experiments/ClimaEarth/user_io/postprocessing.jl +++ b/experiments/ClimaEarth/user_io/postprocessing.jl @@ -3,119 +3,3 @@ include("debug_plots.jl") include("diagnostics_plots.jl") include("../leaderboard/leaderboard.jl") include("../leaderboard/test_rmses.jl") - -""" - postprocess_sim(cs, postprocessing_vars) - -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(cs, postprocessing_vars) - (; conservation_softfail, rmse_check) = postprocessing_vars - (; - 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, 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, artifacts_dir) - - # If we have enough data (in time, but also enough variables), plot the leaderboard. - # We need pressure to compute the leaderboard. - simdir = CAN.SimDir(atmos_output_dir) - if !isempty(simdir) - pressure_in_output = "pfull" in CAN.available_vars(simdir) - 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 = artifacts_dir - compute_leaderboard(leaderboard_base_path, atmos_output_dir, 3) - rmse_check && test_rmse_thresholds(atmos_output_dir, 3) - pressure_in_output && - compute_pfull_leaderboard(leaderboard_base_path, atmos_output_dir, 3) - end - end - - # Perform conservation checks if they exist - if !isnothing(cs.conservation_checks) - @info "Conservation Check Plots" - plot_global_conservation( - cs.conservation_checks.energy, - cs, - conservation_softfail, - 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(artifacts_dir, "total_water_bucket.png"), - figname2 = joinpath(artifacts_dir, "total_water_log_bucket.png"), - ) - end - - # Close all diagnostics file writers - !isnothing(cs.diags_handler) && - map(diag -> close(diag.output_writer), cs.diags_handler.scheduled_diagnostics) -end - -""" - simulated_years_per_day(cs, walltime) - -Compute the simulated years per walltime day for the given coupled simulation `cs`, assuming -that the simulation took `walltime`. -""" -function simulated_years_per_day(cs, walltime) - simulated_seconds_per_second = float(cs.tspan[end] - cs.tspan[begin]) / walltime - return simulated_seconds_per_second / 365.25 -end - -""" - walltime_per_coupling_step(cs, walltime) - -Compute the average walltime needed to take one step for the given coupled simulation `cs`, -assuming that the simulation took `walltime`. The result is in seconds. -""" -function walltime_per_coupling_step(cs, walltime) - n_coupling_steps = (cs.tspan[end] - cs.tspan[begin]) / cs.Δt_cpl - return walltime / n_coupling_steps -end - - -""" - save_sypd_walltime_to_disk(cs, walltime) - -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.dir_paths.artifacts_dir, "sypd.txt"), "w") do sypd_filename - println(sypd_filename, "$sypd") - end - - open( - 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 - end - return nothing -end diff --git a/src/Checkpointer.jl b/src/Checkpointer.jl index 34a827bed4..b13921b076 100644 --- a/src/Checkpointer.jl +++ b/src/Checkpointer.jl @@ -7,14 +7,16 @@ module Checkpointer import ClimaComms import ClimaCore as CC +import CC.DataLayouts: AbstractData import ClimaUtilities.Utils: sort_by_creation_time import ClimaUtilities.TimeManager: ITime, seconds -import ..Interfacer +import ClimaUtilities.TimeVaryingInputs: AbstractTimeVaryingInput import Dates - import JLD2 +import StaticArrays +import ..Interfacer -export get_model_prog_state, checkpoint_model_state, checkpoint_sims +export get_model_prog_state, checkpoint_model_state, checkpoint_sims, restore! """ get_model_prog_state(sim::Interfacer.ComponentModelSimulation) @@ -295,4 +297,111 @@ function remove_checkpoint(prev_checkpoint_file, prev_checkpoint_t, comms_ctx) return nothing end +""" + restore!(v1, v2, comms_ctx; ignore) + +Recursively traverse `v1` and `v2`, setting each field of `v1` with the +corresponding field in `v2`. In this, ignore all the properties that have name +within the `ignore` iterable. + +`ignore` is useful when there are stateful properties, such as live pointers. +""" +function restore!(v1::T1, v2::T2, comms_ctx; name = "", ignore) where {T1, T2} + # We pick fieldnames(T2) because v2 tend to be simpler (Array as opposed + # to CuArray) + fields = filter(x -> !(x in ignore), fieldnames(T2)) + if isempty(fields) + v1 == v2 || error("$v1 != $v2") + else + # Recursive case + for p in fields + restore!( + getfield(v1, p), + getfield(v2, p), + comms_ctx; + name = "$(name).$(p)", + ignore, + ) + end + end + return nothing +end + +# Ignoring certain types that don't need to be restored +# UnionAll and DataType are infinitely recursive, so we also ignore those +function restore!( + v1::Union{ + AbstractTimeVaryingInput, + ClimaComms.AbstractCommsContext, + ClimaComms.AbstractDevice, + UnionAll, + DataType, + }, + v2::Union{ + AbstractTimeVaryingInput, + ClimaComms.AbstractCommsContext, + ClimaComms.AbstractDevice, + UnionAll, + DataType, + }, + _comms_ctx; + name, + ignore, +) + return nothing +end + +function restore!( + v1::Union{AbstractData, AbstractArray}, + v2::Union{AbstractData, AbstractArray}, + comms_ctx; + name, + ignore, +) + ArrayType = + parent(v1) isa Array ? Array : ClimaComms.array_type(ClimaComms.device(comms_ctx)) + moved_to_device = ArrayType(parent(v2)) + + parent(v1) .= moved_to_device + return nothing +end + +function restore!( + v1::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol}, + v2::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol}, + comms_ctx; + name, + ignore, +) + v1 == v2 || error("$name is a immutable but it inconsistent ($(v1) != $(v2))") + return nothing +end + +function restore!(v1::Dict, v2::Dict, comms_ctx; name, ignore) + # RRTGMP has some internal dictionaries + v1 == v2 || error("$name is inconsistent") + return nothing +end + +""" + restore!(v1::T1, v2::T2, comms_ctx; name = "", ignore) where {T1 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond}, T2 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond}} + +Special case for time-related types to allow different timestamps during restore. +""" +function restore!( + v1::T1, + v2::T2, + comms_ctx; + name, + ignore, +) where { + T1 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond}, + T2 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond}, +} + if v1 != v2 + @warn "Time value differs in restart" field = name original = v2 new = v1 + end + return nothing +end + end # module diff --git a/src/ClimaCoupler.jl b/src/ClimaCoupler.jl index 7dc7fb5d8f..5703bc29dd 100644 --- a/src/ClimaCoupler.jl +++ b/src/ClimaCoupler.jl @@ -1,7 +1,7 @@ """ ClimaCoupler -Module for atmos-ocean-land coupled simulations. +Module for atmos-ocean-land-ice coupled simulations. """ module ClimaCoupler @@ -12,5 +12,12 @@ include("ConservationChecker.jl") include("FluxCalculator.jl") include("FieldExchanger.jl") include("Checkpointer.jl") +include("Input.jl") +include("Postprocessor/Postprocessor.jl") +include("SimCoordinator.jl") + +# Re-export commonly used functions for convenience +export run!, step! # from SimCoordinator +export postprocess # from Postprocessor end diff --git a/experiments/ClimaEarth/user_io/arg_parsing.jl b/src/Input.jl similarity index 50% rename from experiments/ClimaEarth/user_io/arg_parsing.jl rename to src/Input.jl index 184fb28b63..d508f45b54 100644 --- a/experiments/ClimaEarth/user_io/arg_parsing.jl +++ b/src/Input.jl @@ -1,14 +1,249 @@ -import YAML +""" + Input + +This module handles parsing of configuration files, command-line arguments, and +other user inputs for setting up coupled simulations. +""" +module Input -mode_name_dict = Dict( - "amip" => AMIPMode, - "cmip" => CMIPMode, - "slabplanet" => SlabplanetMode, - "slabplanet_aqua" => SlabplanetAquaMode, - "slabplanet_terra" => SlabplanetTerraMode, - "subseasonal" => SubseasonalMode, +import ArgParse +import YAML +import Dates +import ClimaAtmos as CA +import ..Interfacer +import ..Utilities +import ClimaUtilities.TimeManager: ITime +import ClimaCoupler + +export argparse_settings, parse_commandline, get_coupler_config_dict, get_coupler_args, parse_component_dts! + +# Mode name dictionary for converting strings to types +const mode_name_dict = Dict( + "amip" => Interfacer.AMIPMode, + "cmip" => Interfacer.CMIPMode, + "slabplanet" => Interfacer.SlabplanetMode, + "slabplanet_aqua" => Interfacer.SlabplanetAquaMode, + "slabplanet_terra" => Interfacer.SlabplanetTerraMode, + "subseasonal" => Interfacer.SubseasonalMode, ) +""" + argparse_settings() + +Set up command-line argument parsing settings for ClimaCoupler. +""" +function argparse_settings() + s = ArgParse.ArgParseSettings() + ArgParse.@add_arg_table! s begin + ### ClimaCoupler flags + # Simulation-identifying information + "--config_file" + help = "A yaml file used to set the configuration of the coupled model [\"config/ci_configs/amip_default.yml\" (default)]" + arg_type = String + default = joinpath(pkgdir(ClimaCoupler), "config/ci_configs/amip_default.yml") + "--job_id" + help = "A unique identifier for this run, defaults to the config file name" + arg_type = String + default = nothing + "--print_config_dict" + help = "Boolean flag indicating whether to print the final configuration dictionary [`true` (default), `false`]" + arg_type = Bool + default = true + "--mode_name" + help = "Mode of coupled simulation. [`cmip`, `amip` (default), `subseasonal`, `slabplanet`, `slabplanet_aqua`, `slabplanet_terra`]" + arg_type = String + default = "amip" + "--coupler_toml" + help = "An optional list of paths to toml files used to overwrite the default model parameters." + arg_type = Vector{String} + default = [] + # Computational simulation setup information + "--unique_seed" + help = "Boolean flag indicating whether to set the random number seed to a unique value [`false` (default), `true`]" + arg_type = Bool + default = false + "--FLOAT_TYPE" + help = "Floating point precision [`Float64` (default), `Float32`]" + arg_type = String + default = "Float64" + "--device" + help = "Device type to use [\"auto\" (default), \"CPUSingleThreaded\", \"CPUMultiThreaded\", \"CUDADevice\"]" + arg_type = String + default = "auto" + # Time information + "--use_itime" + help = "Boolean flag indicating whether to use ITime (integer time) or not (will use Float64) [`true` (default), `false`]" + arg_type = Bool + default = true + "--t_end" + help = "End time of the simulation [\"800secs\"; allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + default = "800secs" + "--t_start" + help = "Start time of the simulation [\"0secs\" (default); allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + default = "0secs" + "--start_date" + help = "Start date of the simulation, in format \"YYYYMMDD\" [\"20100101\" (default)]" + arg_type = String + default = "20000101" + "--dt_cpl" + help = "Coupling time step in seconds [400 (default); allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + default = "400secs" + "--dt" + help = "Component model time step [allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + default = "400secs" + "--dt_atmos" + help = "Atmos simulation time step (alternative to `dt`; no default) [allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + "--dt_land" + help = "Land simulation time step (alternative to `dt`; no default) [allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + "--dt_ocean" + help = "Ocean simulation time step (alternative to `dt`; no default) [allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + "--dt_seaice" + help = "Sea ice simulation time step (alternative to `dt`; no default) [allowed formats: \"Nsecs\", \"Nmins\", \"Nhours\", \"Ndays\", \"Inf\"]" + arg_type = String + "--checkpoint_dt" + help = "Time interval for checkpointing [\"90days\" (default)]" + arg_type = String + default = "90days" + # Space information + "--h_elem" + help = "Number of horizontal elements to use for the boundary space [16 (default)]" + arg_type = Int + default = 16 + "--share_surface_space" + help = "Boolean flag indicating whether to share the surface space between the surface models, atmosphere, and boundary [`true` (default), `false`]" + arg_type = Bool + default = true + # Restart information + "--detect_restart_files" + help = "Boolean flag indicating whether to automatically use restart files if available [`false` (default), `true`]" + arg_type = Bool + default = false + "--restart_dir" + help = "Directory containing restart files" + arg_type = String + default = nothing + "--restart_t" + help = "Time in seconds rounded to the nearest index to use at `t_start` for restarted simulation [nothing (default)]" + arg_type = Int + default = nothing + "--restart_cache" + help = "Boolean flag indicating whether to read the cache from the restart file if available [`true` (default), `false`]" + arg_type = Bool + default = true + "--save_cache" + help = "Boolean flag indicating whether to save the state and cache or only the state when checkpointing [`true` (default), `false`]" + arg_type = Bool + default = true + # Diagnostics information + "--use_coupler_diagnostics" + help = "Boolean flag indicating whether to compute and output coupler diagnostics [`true` (default), `false`]" + arg_type = Bool + default = true + # Physical simulation information + "--evolving_ocean" + help = "Boolean flag indicating whether to use a dynamic slab ocean model, as opposed to constant surface temperatures [`true` (default), `false`]" + arg_type = Bool + default = true + # Conservation and RMSE check information + "--energy_check" + help = "Boolean flag indicating whether to check energy conservation [`false` (default), `true`]" + arg_type = Bool + default = false + "--conservation_softfail" + help = "Boolean flag indicating whether to soft fail on conservation errors [`false` (default), `true`]" + arg_type = Bool + default = false + "--rmse_check" + help = "Boolean flag indicating whether to check RMSE of some physical fields [`false` (default), `true`]" + arg_type = Bool + default = false + # Output information + "--coupler_output_dir" + help = "Directory to save output files. Note that TempestRemap fails if interactive and paths are too long. [\"experiments/ClimaEarth/output\" (default)]" + arg_type = String + default = "experiments/ClimaEarth/output" + # ClimaAtmos specific + "--surface_setup" + help = "Triggers ClimaAtmos into the coupled mode [`PrescribedSurface` (default), `DefaultMoninObukhov\"]" # retained here for standalone Atmos benchmarks + arg_type = String + default = "PrescribedSurface" + "--atmos_config_file" + help = "An optional YAML file used to overwrite the default model parameters." + arg_type = String + default = nothing + "--atmos_log_progress" + help = "Use the ClimaAtmos walltime logging callback instead of the default ClimaCoupler one [`false` (default), `true`]" + arg_type = Bool + default = false + "--albedo_model" + help = "Type of albedo model. [`ConstantAlbedo`, `RegressionFunctionAlbedo`, `CouplerAlbedo` (default)]" + arg_type = String + default = "CouplerAlbedo" + "--extra_atmos_diagnostics" + help = "List of dictionaries containing information about additional atmosphere diagnostics to output [nothing (default)]" + arg_type = Vector{Dict{Any, Any}} + default = [] + ### ClimaLand specific + "--land_model" + help = "Land model to use. [`bucket` (default), `integrated`]" + arg_type = String + default = "bucket" + "--land_temperature_anomaly" + help = "Type of temperature anomaly for land model. [`amip`, `aquaplanet` (default), `nothing`]" + arg_type = String + default = "aquaplanet" + "--use_land_diagnostics" + help = "Boolean flag indicating whether to compute and output land model diagnostics [`true` (default), `false`]" + arg_type = Bool + default = true + "--land_spun_up_ic" + help = "Boolean flag to indicate whether to use integrated land initial conditions from spun up state [`true` (default), `false`]" + arg_type = Bool + default = true + # BucketModel specific + "--bucket_albedo_type" + help = "Access bucket surface albedo information from data file. [`map_static` (default), `function`, `map_temporal`]" + arg_type = String + default = "map_static" # to be replaced by land config file, when available + "--bucket_initial_condition" + help = "A file path for a NetCDF file (read documentation about requirements)" + arg_type = String + default = "" + "--era5_initial_condition_dir" + help = "Directory containing ERA5 initial condition files (subseasonal mode). Filenames inferred from start_date [none (default)]. Generated with `https://github.com/CliMA/WeatherQuest`" + arg_type = String + default = nothing + # Ice model specific + "--ice_model" + help = "Sea ice model to use. [`prescribed` (default), `clima_seaice`]" + arg_type = String + default = "prescribed" + end + return s +end + +""" + parse_commandline(settings) + +Parse command-line arguments using the provided ArgParse settings. +This is intended to be used with the `argparse_settings` function, e.g: +`parsed_args = Input. parse_commandline(Input.argparse_settings())`. + +# Arguments +- `settings`: An `ArgParse.ArgParseSettings` object containing the command-line arguments to parse + +# Returns +- A dictionary of parsed command-line arguments +""" +parse_commandline(settings) = ArgParse.parse_args(ARGS, settings) + """ get_coupler_config_dict(config_file) @@ -21,7 +256,7 @@ it with the coupler configuration. The order of priority for overwriting configuration options from lowest to highest is: 1. ClimaAtmos defaults - 2. ClimaCoupler defaults (defined in `experiments/ClimaEarth/cli_options.jl`) + 2. ClimaCoupler defaults (defined in Input module) 3. Command line arguments provided to ClimaCoupler 4. ClimaAtmos configuration file (if specified in coupler config file) 5. ClimaCoupler configuration file @@ -213,8 +448,6 @@ function get_coupler_args(config_dict::Dict) ) end -### Helper functions used in argument parsing ### - """ get_diag_period() @@ -307,3 +540,5 @@ function parse_component_dts!(config_dict) config_dict["component_dt_dict"] = component_dt_dict return nothing end + +end # module diff --git a/src/Postprocessor/Postprocessor.jl b/src/Postprocessor/Postprocessor.jl new file mode 100644 index 0000000000..d4875223dd --- /dev/null +++ b/src/Postprocessor/Postprocessor.jl @@ -0,0 +1,246 @@ +""" + Postprocessor + +This module handles postprocessing of coupled simulations, including generating +plots, checking conservation, and other diagnostics. +""" +module Postprocessor + +import ClimaComms +import ClimaDiagnostics as CD +import ClimaAnalysis as CAN +import ..Interfacer +import ..ConservationChecker +import Dates + +export postprocess, postprocess_sim, simulated_years_per_day, walltime_per_coupling_step, save_sypd_walltime_to_disk +export make_diagnostics_plots, make_ocean_diagnostics_plots, debug, plot_global_conservation +export compute_leaderboard, compute_pfull_leaderboard, test_rmse_thresholds +export coupler_diagnostics_setup + +""" + postprocess(cs; conservation_softfail = false, rmse_check = false) + +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::Interfacer.CoupledSimulation; conservation_softfail = false, rmse_check = false) + if ClimaComms.iamroot(ClimaComms.context(cs)) && !isnothing(cs.diags_handler) + postprocessing_vars = (; conservation_softfail, rmse_check) + postprocess_sim(cs, postprocessing_vars) + end + return nothing +end + +""" + postprocess_sim(cs, postprocessing_vars) + +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. + +Note: This function depends on ClimaEarth-specific helper functions that are +expected to be available in the calling context (e.g., from ClimaEarth experiments). +""" +function postprocess_sim(cs::Interfacer.CoupledSimulation, postprocessing_vars) + (; conservation_softfail, rmse_check) = postprocessing_vars + (; + coupler_output_dir, + atmos_output_dir, + land_output_dir, + ocean_output_dir, + artifacts_dir, + ) = cs.dir_paths + + # Plot generic diagnostics (ClimaEarth-specific functions) + @info "Plotting diagnostics for coupler, atmos, land, and 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) + if ClimaComms.context(cs) isa ClimaComms.SingletonCommsContext + debug(cs, artifacts_dir) + end + + # If we have enough data (in time, but also enough variables), plot the leaderboard. + # We need pressure to compute the leaderboard. + simdir = CAN.SimDir(atmos_output_dir) + if !isempty(simdir) + pressure_in_output = "pfull" in CAN.available_vars(simdir) + first_var = first(CAN.available_vars(simdir)) + times = CAN.times(CAN.get(simdir; short_name = first_var)) + t_end = times[end] + if t_end > 84600 * 31 * 3 # 3 months for spin up + 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 && + compute_pfull_leaderboard(leaderboard_base_path, atmos_output_dir, 3) + end + end + + # Perform conservation checks if they exist + if !isnothing(cs.conservation_checks) + @info "Conservation Check Plots" + plot_global_conservation( + cs.conservation_checks.energy, + cs, + conservation_softfail, + 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(artifacts_dir, "total_water_bucket.png"), + figname2 = joinpath(artifacts_dir, "total_water_log_bucket.png"), + ) + end + + # Close all diagnostics file writers + !isnothing(cs.diags_handler) && + map(diag -> close(diag.output_writer), cs.diags_handler.scheduled_diagnostics) + return nothing +end + +""" + simulated_years_per_day(cs, walltime) + +Compute the simulated years per walltime day for the given coupled simulation `cs`, assuming +that the simulation took `walltime`. +""" +function simulated_years_per_day(cs::Interfacer.CoupledSimulation, walltime) + simulated_seconds_per_second = float(cs.tspan[end] - cs.tspan[begin]) / walltime + return simulated_seconds_per_second / 365.25 +end + +""" + walltime_per_coupling_step(cs, walltime) + +Compute the average walltime needed to take one step for the given coupled simulation `cs`, +assuming that the simulation took `walltime`. The result is in seconds. +""" +function walltime_per_coupling_step(cs::Interfacer.CoupledSimulation, walltime) + n_coupling_steps = (cs.tspan[end] - cs.tspan[begin]) / cs.Δt_cpl + return walltime / n_coupling_steps +end + +""" + save_sypd_walltime_to_disk(cs, walltime) + +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::Interfacer.CoupledSimulation, 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.dir_paths.artifacts_dir, "sypd.txt"), "w") do sypd_filename + println(sypd_filename, "$sypd") + end + + open( + 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 + end + return nothing +end + +""" + CD.orchestrate_diagnostics(cs::CoupledSimulation) + +Compute and output coupled diagnostics. +""" +function CD.orchestrate_diagnostics(cs::Interfacer.CoupledSimulation) + ## wrap the current CoupledSimulation fields and time in a NamedTuple to match the ClimaDiagnostics interface + cs_nt = (; u = cs.fields, p = nothing, t = cs.t[], step = round(cs.t[] / cs.Δt_cpl)) + !isnothing(cs.diags_handler) && CD.orchestrate_diagnostics(cs_nt, cs.diags_handler) + return nothing +end + +""" + coupler_diagnostics_setup(fields, output_dir, start_date, t_start, diagnostics_dt, coupled_dt) + +Set up the default diagnostics for an AMIP simulation, using ClimaDiagnostics. +The diagnostics are saved to NetCDF files. Currently, this just includes a +diagnostic for turbulent energy fluxes. + +Return a DiagnosticsHandler object to coordinate the diagnostics. +""" +function coupler_diagnostics_setup( + fields, + output_dir, + start_date, + t_start, + diagnostics_dt, + coupled_dt, +) + # Create schedules and writer + schedule_everystep = CD.Schedules.EveryStepSchedule() + schedule_calendar_dt = CD.Schedules.EveryCalendarDtSchedule(diagnostics_dt; start_date) + netcdf_writer = CD.Writers.NetCDFWriter(axes(fields.F_sh), output_dir) + + # Create the diagnostic for turbulent energy fluxes + F_turb_energy_diag = CD.DiagnosticVariable(; + short_name = "F_turb_energy", + long_name = "Turbulent energy fluxes", + standard_name = "F_turb_energy", + units = "W m^-2", + comments = "Turbulent energy fluxes are calculated as the sum of sensible and latent heat fluxes, + weighted by surface simulation area.", + compute! = (out, state, cache, time) -> begin + if isnothing(out) + return state.F_sh .+ state.F_lh + else + out .= state.F_sh .+ state.F_lh + end + end, + ) + + # Schedule the turbulent energy fluxes to save at every step, and output at the frequency calculated above + F_turb_energy_diag_sched = CD.ScheduledDiagnostic( + variable = F_turb_energy_diag, + output_writer = netcdf_writer, + reduction_time_func = (+), + compute_schedule_func = schedule_everystep, + output_schedule_func = schedule_calendar_dt, + pre_output_hook! = CD.average_pre_output_hook!, + ) + + # Create the diagnostics handler containing the scheduled diagnostics + scheduled_diags = [F_turb_energy_diag_sched] + diags_handler = + CD.DiagnosticsHandler(scheduled_diags, fields, nothing, t_start, dt = coupled_dt) + return diags_handler +end + +include("diagnostics_plots.jl") +include("debug_plots.jl") +include("benchmarks.jl") +include("leaderboard/leaderboard.jl") +include("leaderboard/test_rmses.jl") + +end # module diff --git a/src/Postprocessor/benchmarks.jl b/src/Postprocessor/benchmarks.jl new file mode 100644 index 0000000000..b7bec75991 --- /dev/null +++ b/src/Postprocessor/benchmarks.jl @@ -0,0 +1,189 @@ +#= +Our goal here is to output a table displaying some results from benchmark runs +in the coupler. We mainly want to compare performance (via SYPD) between +coupled and atmos-only runs. + +The table should look something like this: +------------------------------------------------------- +| Build ID: XYZ | Horiz. res: ___ | GPU Run [X A100s] | +| | Vert. res: ___ | | +| | dt: ______ | | +------------------------------------------------------- +| | job ID: | $job_id | +| Coupled run | SYPD: | $SYPD | +------------------------------------------------------- +| | job ID: | $job_id | +| Atmos-only | SYPD: | $SYPD | +------------------------------------------------------- +=# + +import ArgParse +import PrettyTables +import ClimaCoupler + +function benchmarks_argparse_settings() + s = ArgParse.ArgParseSettings() + ArgParse.@add_arg_table! s begin + "--gpu_job_id_coupled" + help = "The name of the GPU coupled run we want to compare." + arg_type = String + default = nothing + "--gpu_job_id_coupled_io" + help = "The name of the GPU coupled with IO run we want to compare." + arg_type = String + default = nothing + "--gpu_job_id_atmos" + help = "The name of the GPU atmos-only run we want to compare." + arg_type = String + default = nothing + "--gpu_job_id_atmos_diagedmf" + help = "The name of the GPU atmos-only run we want to compare." + arg_type = String + default = nothing + "--gpu_job_id_coupled_progedmf_coarse" + help = "The name of the coarser GPU coupled with prognostic EDMF + 1M run we want to compare." + arg_type = String + default = nothing + "--gpu_job_id_coupled_progedmf_fine" + help = "The name of the finer GPU coupled with prognostic EDMF + 1M run we want to compare." + arg_type = String + default = nothing + "--coupler_output_dir" + help = "Directory to save output files." + arg_type = String + default = "experiments/ClimaEarth/output" + "--build_id" + help = "The build ID of the pipeline running this script." + arg_type = String + default = nothing + end + return s +end + +""" + get_run_info(parsed_args, run_type) + +Use the input `parsed_args` to get the job ID and artifacts directories for +the GPU run of the given `run_type`. + +`run_type` must be one of "coupled", "coupled_io", "atmos", or "atmos_diagedmf". +""" +function get_run_info(parsed_args, run_type) + # Read in GPU job ID info from command line + if run_type == "coupled" + gpu_job_id = parsed_args["gpu_job_id_coupled"] + mode_name = "amip" + elseif run_type == "coupled_io" + gpu_job_id = parsed_args["gpu_job_id_coupled_io"] + mode_name = "amip" + elseif run_type == "atmos_diagedmf" + gpu_job_id = parsed_args["gpu_job_id_atmos_diagedmf"] + mode_name = "climaatmos" + elseif run_type == "atmos" + gpu_job_id = parsed_args["gpu_job_id_atmos"] + mode_name = "climaatmos" + elseif run_type == "coupled_progedmf_coarse" + gpu_job_id = parsed_args["gpu_job_id_coupled_progedmf_coarse"] + mode_name = "amip" + elseif run_type == "coupled_progedmf_fine" + gpu_job_id = parsed_args["gpu_job_id_coupled_progedmf_fine"] + mode_name = "amip" + else + error("Invalid run type: $run_type") + end + + # Verify that the user has provided the necessary job IDs + isnothing(gpu_job_id) && error("Must pass GPU coupled run name to compare it.") + + # Construct GPU artifacts directory + gpu_artifacts_dir = joinpath(output_dir, gpu_job_id, "artifacts") + + return (gpu_job_id, gpu_artifacts_dir) +end + +""" + get_run_data(artifacts_dir) + +Read in run data from artifacts directory, currently SYPD. +""" +function get_run_data(artifacts_dir) + # Read in SYPD info + sypd = open(joinpath(artifacts_dir, "sypd.txt"), "r") do sypd_file + round(parse(Float64, read(sypd_file, String)), digits = 4) + end + + return sypd +end + +""" + append_table_data(table_data, setup_id, gpu_job_id, gpu_artifacts_dir) + +Append data for a given setup to the table data. +""" +function append_table_data(table_data, setup_id, gpu_job_id, gpu_artifacts_dir) + # Get SYPD info for the GPU run + gpu_sypd = get_run_data(gpu_artifacts_dir) + + # Create rows containing data for these runs + new_table_data = [ + [setup_id "job ID:" gpu_job_id] + ["" "SYPD:" gpu_sypd] + ] + return vcat(table_data, new_table_data) +end +# Read in command line arguments +parsed_args = ArgParse.parse_args(ARGS, benchmarks_argparse_settings()) +output_dir = parsed_args["coupler_output_dir"] + +# Access buildkite pipeline ID (from `BUILDKITE_GITHUB_DEPLOYMENT_ID` variable) +build_id = parsed_args["build_id"] +if !isnothing(build_id) + build_id_str = "Build ID: $build_id" +else + build_id_str = "Build ID: N/A" +end + +# Read in run info for each of the cases we want to compare +run_info_coupled_progedmf_coarse = get_run_info(parsed_args, "coupled_progedmf_coarse") +run_info_coupled_progedmf_fine = get_run_info(parsed_args, "coupled_progedmf_fine") +run_info_coupled_io = get_run_info(parsed_args, "coupled_io") +run_info_coupled = get_run_info(parsed_args, "coupled") +run_info_atmos_diagedmf = get_run_info(parsed_args, "atmos_diagedmf") +run_info_atmos = get_run_info(parsed_args, "atmos") + +# Set up info for PrettyTables.jl +headers = [build_id_str, "Horiz. res.: 30 elems", "GPU Run [2 A100s]"] +data = [ + ["" "Vert. res.: 63 levels" ""] + ["" "dt: 120secs" ""] +] + +# Append data to the table for each of the cases we want to compare +data = append_table_data( + data, + "Coupled with progedmf + 1M (16 helem)", + run_info_coupled_progedmf_coarse..., +) +data = append_table_data( + data, + "Coupled with progedmf + 1M (30 helem)", + run_info_coupled_progedmf_fine..., +) +data = append_table_data(data, "Coupled with diag. EDMF + IO", run_info_coupled_io...) +data = append_table_data(data, "Coupled with diag. EDMF", run_info_coupled...) +data = append_table_data(data, "Atmos with diag. EDMF", run_info_atmos_diagedmf...) +data = append_table_data(data, "Atmos without diag. EDMF", run_info_atmos...) + +# Output table to file, note that this must match the slack upload command in the pipeline.yml file +table_output_dir = joinpath(output_dir, "compare_amip_climaatmos_amip_diagedmf") +mkpath(table_output_dir) +table_path = joinpath(table_output_dir, "table.txt") +open(table_path, "w") do f + # Output the table, including lines before and after the header + PrettyTables.pretty_table( + f, + data, + header = headers, + hlines = [0, 3, 5, 7, 9, 11, 13, 15], + ) +end diff --git a/src/Postprocessor/debug_plots.jl b/src/Postprocessor/debug_plots.jl new file mode 100644 index 0000000000..4315526376 --- /dev/null +++ b/src/Postprocessor/debug_plots.jl @@ -0,0 +1,298 @@ +import Printf +import ClimaCore as CC +import Makie +import ClimaCoreMakie +import CairoMakie +import ..Interfacer +import ..ConservationChecker +import ClimaAtmos as CA +import Oceananigans as OC +import StaticArrays + +""" + plot_global_conservation( + cc::AbstractConservationCheck, + coupler_sim::Interfacer.CoupledSimulation, + softfail = false; + figname1 = "total.png", + figname2 = "total_log.png", + ) + +Creates two plots of the globally integrated quantity (energy, ``\\rho e``): +1. global quantity of each model component as a function of time, +relative to the initial value; +2. fractional change in the sum of all components over time on a log scale. +""" +function plot_global_conservation( + cc::ConservationChecker.AbstractConservationCheck, + coupler_sim::Interfacer.CoupledSimulation, + softfail = false; + figname1 = "total.png", + figname2 = "total_log.png", +) + + model_sims = coupler_sim.model_sims + ccs = cc.sums + + days = collect(1:length(ccs[1])) * float(coupler_sim.Δt_cpl) / 86400 + + # evolution of energy of each component relative to initial value + total = ccs.total # total + + var_name = nameof(cc) + cum_total = [0.0] + f = Makie.Figure() + ax = Makie.Axis(f[1, 1], xlabel = "time [days]", ylabel = "$var_name: (t) - (t=0)") + Makie.lines!(ax, days, total .- total[1], label = "total"; linewidth = 3) + for sim in model_sims + sim_name = nameof(sim) + global_field = getproperty(ccs, Symbol(sim_name)) + diff_global_field = (global_field .- global_field[1]) + Makie.lines!(ax, days, diff_global_field[1:length(days)], label = sim_name) + cum_total .+= abs.(global_field[end]) + end + if cc isa ConservationChecker.EnergyConservationCheck + global_field = ccs.toa_net_source + diff_global_field = (global_field .- global_field[1]) + Makie.lines!(ax, days, diff_global_field[1:length(days)], label = "toa_net") + cum_total .+= abs.(global_field[end]) + end + Makie.axislegend(ax, position = :lb) + Makie.save(figname1, f) + + # use the cumulative global sum at the final time step as a reference for the error calculation + rse = abs.((total .- total[1]) ./ cum_total) + l_rse = log.(rse) + # evolution of log error of total + lp = Makie.lines(days, l_rse, label = "rs error") + lp.axis.xlabel = "time [days]" + lp.axis.ylabel = "log( |x(t) - x(t=0)| / Σx(t=T) )" + l_rse_valid = filter(x -> !isinf(x) && !isnan(x), l_rse) + if !isempty(l_rse_valid) + y_min = minimum(l_rse_valid) + y_max = maximum(l_rse_valid) + if y_min != y_max + Makie.ylims!(y_min, y_max) + else + # If all values are the same, add a small padding to avoid Makie error + padding = max(abs(y_min) * 0.01, 0.1) + Makie.ylims!(y_min - padding, y_max + padding) + end + end + Makie.axislegend(position = :lt) + Makie.save(figname2, lp) + + # check that the relative error is small (TODO: reduce this to sqrt(eps(FT))) + if !softfail + @info typeof(cc) + @info rse[end] + @assert rse[end] < 0.035 + end +end + + +# plotting functions for the coupled simulation +""" + debug(cs::Interfacer.CoupledSimulation, dir = "debug", cs_fields_ref = nothing) + +Plot the fields of a coupled simulation and save plots to a directory. +""" +function debug(cs::Interfacer.CoupledSimulation, dir = "debug", cs_fields_ref = nothing) + isdir(dir) || mkpath(dir) + @info "plotting debug in $dir" + for sim in cs.model_sims + debug(sim, dir) + end + debug(cs.fields, dir, cs_fields_ref) +end + +""" + debug(cs_fields::CC.Fields.Field, dir, cs_fields_ref = nothing) + +Plot useful coupler fields (in `field_names`) and save plots to a directory. + +If `cs_fields_ref` is provided (e.g., using a copy of cs.fields from the initialization), +plot the anomalies of the fields with respect to `cs_fields_ref`. +""" +function debug(cs_fields::CC.Fields.Field, dir, cs_fields_ref = nothing) + field_names = propertynames(cs_fields) + fig = Makie.Figure(size = (1500, 800)) + min_square_len = ceil(Int, sqrt(length(field_names))) + for i in 1:min_square_len, j in 1:min_square_len + field_index = (i - 1) * min_square_len + j + if field_index <= length(field_names) + field_name = field_names[field_index] + field = getproperty(cs_fields, field_name) + _heatmap_cc_field!(fig, field, i, j, field_name) + end + end + Makie.save(joinpath(dir, "debug_coupler.png"), fig) + + # plot anomalies if a reference cs.fields, `cs_fields_ref`, are provided + if !isnothing(cs_fields_ref) + for i in 1:min_square_len, j in 1:min_square_len + field_index = (i - 1) * min_square_len + j + if field_index <= length(field_names) + field_name = field_names[field_index] + field = getproperty(cs_fields, field_name) + _heatmap_cc_field!(fig, field, i, j, field_name) + end + end + Makie.save(joinpath(dir, "debug_coupler_anomalies.png"), fig) + end +end + +""" + debug(sim::Interfacer.ComponentModelSimulation, dir) + +Plot the fields of a component model simulation and save plots to a directory. +""" +function debug(sim::Interfacer.ComponentModelSimulation, dir) + field_names = plot_field_names(sim) + fig = Makie.Figure(size = (1500, 800)) + min_square_len = ceil(Int, sqrt(length(field_names))) + for i in 1:min_square_len, j in 1:min_square_len + field_index = (i - 1) * min_square_len + j + if field_index <= length(field_names) + field_name = field_names[field_index] + field = Interfacer.get_field(sim, Val(field_name)) + # If field is a ClimaCore field, then _heatmap_cc_field! will add a + # title to the axis + title = + field isa CC.Fields.Field ? "" : string(field_name) * print_extrema(field) + ax = Makie.Axis(fig[i, j * 2 - 1]; title) + if field isa OC.Field || field isa OC.AbstractOperations.AbstractOperation + if field isa OC.AbstractOperations.AbstractOperation + field = OC.Field(field) + OC.compute!(field) + end + grid = field.grid + hm = CairoMakie.heatmap!(ax, view(field, :, :, grid.Nz)) + Makie.Colorbar(fig[i, j * 2], hm) + elseif field isa CC.Fields.Field + _heatmap_cc_field!(fig, field, i, j, field_name) + elseif field isa AbstractArray + lin = Makie.lines!(ax, Array(field)) + end + end + end + Makie.save(joinpath(dir, "debug_$(nameof(sim)).png"), fig) +end + +""" + _heatmap_cc_field!(fig, field::CC.Fields.Field, i, j, field_name) + +Helper function to plot a heatmap of a ClimaCore field in the given figure at position (i, j). +If the field is constant, skip plotting it to avoid heatmap errors. +""" +function _heatmap_cc_field!(fig, field::CC.Fields.Field, i, j, field_name) + # Copy field onto cpu space if necessary + cpu_field = CC.to_cpu(field) + if cpu_field isa CC.Fields.ExtrudedCubedSphereSpectralElementField3D + cpu_field = CC.Fields.level(cpu_field, 1) + end + + # ClimaCoreMakie doesn't support NaNs/Infs, so we substitute them with 100max + FT = CC.Spaces.undertype(axes(cpu_field)) + isinvalid = (x) -> isnan(x) || isinf(x) + field_valid_min, field_valid_max = + extrema(map(x -> isinvalid(x) ? FT(0) : x, parent(cpu_field))) + map!(x -> isinvalid(x) ? 100field_valid_max : x, parent(cpu_field), parent(cpu_field)) + + # If the values are too small, `isapprox` can't be computed accurately because of floating point precision issues. + is_toosmall = (x) -> log10(abs(x)) < log10(floatmin(Float64)) / 2 + + # If the field is constant, skip plotting it to avoid heatmap errors. + if isapprox(field_valid_min, field_valid_max) || + (is_toosmall(field_valid_min) && is_toosmall(field_valid_max)) + ax = Makie.Axis( + fig[i, j * 2 - 1], + title = string(field_name) * print_extrema(cpu_field), + ) + else + colorrange = (field_valid_min, field_valid_max) + ax = Makie.Axis( + fig[i, j * 2 - 1], + title = string(field_name) * print_extrema(cpu_field), + ) + hm = ClimaCoreMakie.fieldheatmap!(ax, cpu_field, colorrange = colorrange) + Makie.Colorbar(fig[i, j * 2], hm) + end + return nothing +end + +""" + print_extrema(field::CC.Fields.Field) + +Return the minimum and maximum values of a field as a string. +""" +function print_extrema(field::Union{CC.Fields.Field, Vector, StaticArrays.SVector}) + ext_vals = extrema(field) + min = Printf.@sprintf("%.2E", ext_vals[1]) + max = Printf.@sprintf("%.2E", ext_vals[2]) + return " [$min, $max]" +end + +function print_extrema(num::Number) + min = Printf.@sprintf("%.2E", num) + max = Printf.@sprintf("%.2E", num) + return " [$min, $max]" +end + +function print_extrema(operation::OC.AbstractOperations.AbstractOperation) + evaluated_field = OC.Field(operation) + OC.compute!(evaluated_field) + return print_extrema(evaluated_field) +end + +function print_extrema(field::OC.Field) + min = Printf.@sprintf("%.2E", minimum(field)) + max = Printf.@sprintf("%.2E", maximum(field)) + return " [$min, $max]" +end + +# below are additional fields specific to this experiment (ourside of the required coupler fields) that we are interested in plotting for debugging purposes + +# additional ClimaAtmos model debug fields + +# additional BucketSimulation debug fields +Interfacer.get_field(sim::BucketSimulation, ::Val{:σS}) = sim.integrator.u.bucket.σS +Interfacer.get_field(sim::BucketSimulation, ::Val{:Ws}) = sim.integrator.u.bucket.Ws +Interfacer.get_field(sim::BucketSimulation, ::Val{:W}) = sim.integrator.u.bucket.W + +# additional ClimaLand model debug fields +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_water}) = + sim.integrator.u.soil.ϑ_l +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_ice}) = sim.integrator.u.soil.θ_i +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:soil_energy}) = + sim.integrator.u.soil.ρe_int +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:canopy_temp}) = + sim.integrator.u.canopy.energy.T +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:canopy_water}) = + sim.integrator.u.canopy.hydraulics.ϑ_l.:1 +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_energy}) = + sim.integrator.u.snow.U +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_water_equiv}) = + sim.integrator.u.snow.S +Interfacer.get_field(sim::ClimaLandSimulation, ::Val{:snow_liquid_water}) = + sim.integrator.u.snow.S_l + +# currently selected plot fields +plot_field_names(sim::Interfacer.SurfaceModelSimulation) = + (:area_fraction, :surface_temperature) +plot_field_names(sim::ClimaLandSimulation) = ( + :area_fraction, + :surface_direct_albedo, + :surface_diffuse_albedo, + :surface_temperature, + :soil_water, + :soil_ice, + :soil_energy, + :canopy_temp, + :canopy_water, + :snow_energy, + :snow_water_equiv, + :snow_liquid_water, +) +plot_field_names(sim::BucketSimulation) = + (:area_fraction, :surface_temperature, :σS, :Ws, :W) diff --git a/src/Postprocessor/diagnostics_plots.jl b/src/Postprocessor/diagnostics_plots.jl new file mode 100644 index 0000000000..0e643b83f8 --- /dev/null +++ b/src/Postprocessor/diagnostics_plots.jl @@ -0,0 +1,276 @@ +import CairoMakie +import CairoMakie.Makie +import ClimaAnalysis as CAN +using Poppler_jll: pdfunite +import Base.Filesystem +import NCDatasets + +const LARGE_NUM = typemax(Int) +const LAST_SNAP = LARGE_NUM +const FIRST_SNAP = -LARGE_NUM +const BOTTOM_LVL = -LARGE_NUM +const TOP_LVL = LARGE_NUM + +function Makie.get_tickvalues(yticks::Int, ymin, ymax) + return range(ymin, ymax, yticks) +end + +YLINEARSCALE = + Dict(:axis => CAN.Utils.kwargs(dim_on_y = true, yticks = 10, ytickformat = "{:.3e}")) + +long_name(var) = var.attributes["long_name"] +short_name(var) = var.attributes["short_name"] + +""" + make_plots_generic( + file_path::Union{<:AbstractString, Vector{<:AbstractString}}, + plot_path, + vars, + args...; + plot_fn = nothing, + output_name = "summary", + summary_files = String[], + MAX_NUM_COLS = 1, + MAX_NUM_ROWS = min(4, length(vars)), + kwargs..., + ) +Create plots for each variable in `vars` and save them to `plot_path`. The number of plots per +page is determined by `MAX_NUM_COLS` and `MAX_NUM_ROWS`. The `plot_fn` function is used to create the +plots. If `plot_fn` is not provided, a default plotting function is used. The default plotting function +is determined by the keyword arguments `kwargs`. +""" +function make_plots_generic( + file_path::Union{<:AbstractString, Vector{<:AbstractString}}, + plot_path, + vars, + args...; + plot_fn = nothing, + output_name = "summary", + summary_files = String[], + MAX_NUM_COLS = 1, + MAX_NUM_ROWS = min(4, length(vars)), + kwargs..., +) + # When file_path is a Vector with multiple elements, this means that this function is + # being used to produce a comparison plot. In that case, we modify the output name, and + # the number of columns (to match how many simulations we are comparing). + is_comparison = file_path isa Vector + # + # However, we don't want to do this when the vector only contains one element. + if is_comparison && length(file_path) == 1 + # Fallback to the "file_path isa String" case + file_path = file_path[1] + is_comparison = false + end + + if is_comparison + MAX_NUM_COLS = length(file_path) + plot_path = file_path[1] + output_name *= "_comparison" + end + + # Default plotting function needs access to kwargs + if isnothing(plot_fn) + plot_fn = (grid_loc, var) -> CAN.Visualize.plot!(grid_loc, var, args...; kwargs...) + end + + MAX_PLOTS_PER_PAGE = MAX_NUM_ROWS * MAX_NUM_COLS + vars_left_to_plot = length(vars) + + # Define fig, grid, and grid_pos, used below. (Needed for scope) + function makefig() + fig = CairoMakie.Figure(; size = (900, 300 * MAX_NUM_ROWS)) + if is_comparison + for (col, path) in (file_path) + # CairoMakie seems to use this Label to determine the width of the figure. + # Here we normalize the length so that all the columns have the same width. + LABEL_LENGTH = 40 + path = convert(Vector{Float64}, path) + normalized_path = + lpad(path, LABEL_LENGTH + 1, " ")[(end - LABEL_LENGTH):end] + + CairoMakie.Label(fig[0, col], path) + end + end + return fig + end + + # Standardizes grid layout + gridlayout() = + map(1:MAX_PLOTS_PER_PAGE) do i + row = mod(div(i - 1, MAX_NUM_COLS), MAX_NUM_ROWS) + 1 + col = mod(i - 1, MAX_NUM_COLS) + 1 + return fig[row, col] = CairoMakie.GridLayout() + end + + fig = makefig() + grid = gridlayout() + page = 1 + grid_pos = 1 + + for var in vars + if all(isnan, var.data) + @warn "$(short_name(var)) diagnostic is entirely NaN - skipping plot" + vars_left_to_plot -= 1 + continue + end + if minimum(var.data) == maximum(var.data) + @warn "$(short_name(var)) diagnostic is spatially constant - skipping plot" + vars_left_to_plot -= 1 + continue + end + + if grid_pos > MAX_PLOTS_PER_PAGE + fig = makefig() + grid = gridlayout() + grid_pos = 1 + end + + plot_fn(grid[grid_pos], var) + grid_pos += 1 + + # Flush current page + if grid_pos > min(MAX_PLOTS_PER_PAGE, vars_left_to_plot) + file_path = joinpath(plot_path, "$(output_name)_$page.pdf") + CairoMakie.resize_to_layout!(fig) + CairoMakie.save(file_path, fig) + push!(summary_files, file_path) + vars_left_to_plot -= MAX_PLOTS_PER_PAGE + page += 1 + end + end + + # Return early if there are no plots to save + isempty(summary_files) && return nothing + + # Save plots + output_file = joinpath(plot_path, "$(output_name).pdf") + + pdfunite() do unite + run(Cmd([unite, summary_files..., output_file])) + end + + # Cleanup + Filesystem.rm.(summary_files, force = true) + return output_file +end + +function map_comparison(func, simdirs, args) + return vcat([[func(simdir, arg) for simdir in simdirs] for arg in args]...) +end + +""" + make_diagnostics_plots( + output_path::AbstractString, + plot_path::AbstractString; + output_prefix = "", + ) +Create plots for diagnostics. The plots are saved to `plot_path`. +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; + output_prefix = "", +) + simdir = CAN.SimDir(output_path) + 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 = 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)) + periods = CAN.available_periods(simdir; short_name, reduction) + "1d" in periods ? (period = "1d") : (period = first(periods)) + vars[i] = get(simdir; short_name, reduction, period) + end + + # Filter vars into 2D and 3D variable diagnostics vectors + # 3D fields are zonally averaged platted on the lat-z plane + # 2D fields are plotted on the lon-lat plane + vars_3D = + map(var_3D -> CAN.average_lon(var_3D), filter(var -> CAN.has_altitude(var), vars)) + vars_2D = filter(var -> !CAN.has_altitude(var), vars) + + # Generate plots and save in `plot_path` + !isempty(vars_3D) && make_plots_generic( + output_path, + plot_path, + vars_3D, + time = LAST_SNAP, + output_name = output_prefix * "summary_3D", + more_kwargs = YLINEARSCALE, + ) + !isempty(vars_2D) && make_plots_generic( + output_path, + plot_path, + vars_2D, + time = LAST_SNAP, + output_name = output_prefix * "summary_2D", + ) +end + +""" + make_ocean_diagnostics_plots(output_path::AbstractString, plot_path::AbstractString; output_prefix = "") + +Create plots for diagnostics. The plots are saved to `ocean_summary_2D.pdf` in `plot_path`. +This function will plot the following variables, if they have been saved in `output_path`: + - Temperature (`T`) + - Salinity (`S`) + - Zonal velocity (`u`) + - Meridional velocity (`v`) + +For each variable, take the surface level (top level) of the variable +and create a 2D plot. The plots will be saved in a single PDF file. +""" +function make_ocean_diagnostics_plots( + output_path::AbstractString, + plot_path::AbstractString; + output_prefix = "", +) + expected_output_path = joinpath(output_path, "ocean_diagnostics.nc") + isfile(expected_output_path) || return nothing + + # Create an OutputVar for each diagnostic, so we can use ClimaAnalysis to plot + var_names = ["T", "S", "u", "v"] + vars = Array{Union{CAN.OutputVar, Nothing}}(undef, length(var_names)) + for (i, var_name) in enumerate(var_names) + # Create an OutputVar if the variable is available in the output file + try + output_var = CAN.OutputVar(expected_output_path, var_name) + catch e + if e isa NCDatasets.NetCDFError + @warn "NetCDF error with diagnostic $var_name; check that diagnostic is correctly saved" + vars[i] = nothing + continue + else + rethrow(e) + end + end + output_var.attributes["short_name"] = var_name + + # Take the top level (surface) of the variable + output_var = CAN.slice(output_var, z_aac = output_var.dims["z_aac"][1]) + + vars[i] = output_var + end + + # Filter out any variables that are not available + vars = filter(!isnothing, vars) + + # Make plots for each variable, saved in one PDF file + !isempty(vars) && make_plots_generic( + expected_output_path, # file_path + plot_path, + vars, + time = LAST_SNAP, + output_name = output_prefix * "summary_2D", + ) + return nothing +end diff --git a/src/Postprocessor/leaderboard/data_sources.jl b/src/Postprocessor/leaderboard/data_sources.jl new file mode 100644 index 0000000000..47167fff59 --- /dev/null +++ b/src/Postprocessor/leaderboard/data_sources.jl @@ -0,0 +1,452 @@ +import ClimaAnalysis +import ClimaUtilities.ClimaArtifacts: @clima_artifact + +# Tuple of short names for loading simulation and observational data +sim_obs_short_names_rad = [ + ("rsdt", "solar_mon"), + ("rsut", "toa_sw_all_mon"), + ("rlut", "toa_lw_all_mon"), + ("rsutcs", "toa_sw_clr_t_mon"), + ("rlutcs", "toa_lw_clr_t_mon"), + ("rsds", "sfc_sw_down_all_mon"), + ("rsus", "sfc_sw_up_all_mon"), + ("rlds", "sfc_lw_down_all_mon"), + ("rlus", "sfc_lw_up_all_mon"), + ("rsdscs", "sfc_sw_down_clr_t_mon"), + ("rsuscs", "sfc_sw_up_clr_t_mon"), + ("rldscs", "sfc_lw_down_clr_t_mon"), +] + +sim_obs_short_names_sf = [("hfss", "msshf"), ("hfls", "mslhf")] + +""" + get_compare_vars_biases_groups() + +Return an array of arrays of short names. + +This determines which variables are grouped on the bias plots. +""" +function get_compare_vars_biases_groups() + compare_vars_biases_groups = [ + ["pr", "rsdt", "rsut", "rlut"], + ["rsds", "rsus", "rlds", "rlus"], + ["rsutcs", "rlutcs", "rsdscs", "rsuscs", "rldscs"], + ["hfss", "hfls"], + ] + return compare_vars_biases_groups +end + +""" + get_compare_vars_biases_plot_extrema() + +Return a dictionary mapping short names to ranges for the bias plots. This is +used by the function `compute_leaderboard`. + +To add a variable to the leaderboard, add a key-value pair to the dictionary +`compare_vars_biases_plot_extrema` whose key is a short name key is the same +short name in `sim_var_dict` and the value is a tuple, where the first element +is the lower bound and the last element is the upper bound for the bias plots. +""" +function get_compare_vars_biases_plot_extrema() + compare_vars_biases_plot_extrema = Dict( + "pr" => (-5.0, 5.0), + "rsdt" => (-2.0, 2.0), + "rsut" => (-50.0, 50.0), + "rlut" => (-50.0, 50.0), + "rsutcs" => (-20.0, 20.0), + "rlutcs" => (-20.0, 20.0), + "rsds" => (-50.0, 50.0), + "rsus" => (-10.0, 10.0), + "rlds" => (-50.0, 50.0), + "rlus" => (-50.0, 50.0), + "rsdscs" => (-10.0, 10.0), + "rsuscs" => (-10.0, 10.0), + "rldscs" => (-20.0, 20.0), + "hfss" => (-50.0, 50.0), + "hfls" => (-50.0, 50.0), + ) + return compare_vars_biases_plot_extrema +end + +""" + get_sim_var_dict(diagnostics_folder_path) + +Return a dictionary mapping short names to `OutputVar` containing preprocessed +simulation data. This is used by the function `compute_leaderboard`. + +To add a variable for the leaderboard, add a key-value pair to the dictionary +`sim_var_dict` whose key is the short name of the variable and the value is an +anonymous function that returns a `OutputVar`. For each variable, any +preprocessing should be done in the corresponding anonymous function which +includes unit conversion and shifting the dates. + +The variable should have only three dimensions: latitude, longitude, and time. +""" +function get_sim_var_dict(diagnostics_folder_path) + available_short_names = get_short_names_monthly_averages(diagnostics_folder_path) + # Dict for loading in simulation data + sim_var_dict = Dict{String, Any}() + # Add "pr" and the necessary preprocessing + "pr" in available_short_names && ( + sim_var_dict["pr"] = + () -> begin + sim_var = get( + ClimaAnalysis.SimDir(diagnostics_folder_path), + short_name = "pr", + reduction = "average", + period = "1M", + ) + sim_var = ClimaAnalysis.convert_units( + sim_var, + "mm/day", + conversion_function = x -> x .* Float32(-86400), + ) + sim_var = ClimaAnalysis.shift_to_start_of_previous_month(sim_var) + return sim_var + end + ) + + for (short_name, _) in vcat(sim_obs_short_names_rad, sim_obs_short_names_sf) + short_name in available_short_names && ( + sim_var_dict[short_name] = + () -> begin + sim_var = get( + ClimaAnalysis.SimDir(diagnostics_folder_path), + short_name = short_name, + reduction = "average", + period = "1M", + ) + sim_var = ClimaAnalysis.shift_to_start_of_previous_month(sim_var) + return sim_var + end + ) + end + return sim_var_dict +end + +""" + get_obs_var_dict() + +Return a dictionary mapping short names to `OutputVar` containing preprocessed +observational data. This is used by the function `compute_leaderboard`. + +To add a variable for the leaderboard, add a key-value pair to the dictionary +`obs_var_dict` whose key is the short name of the variable and the value is an +anonymous function that returns a `OutputVar`. The function must take in a +start date which is used to align the times in the observational data to match +the simulation data. The short name must be the same as in `sim_var_dict` in the +function `sim_var_dict`. For each variable, any preprocessing is done in the +corresponding anonymous function which includes unit conversion and shifting the +dates. + +The variable should have only three dimensions: latitude, longitude, and time. +""" +function get_obs_var_dict() + obs_var_dict = Dict{String, Any}() + obs_var_dict["pr"] = + (start_date) -> begin + obs_var = ClimaAnalysis.OutputVar( + joinpath(@clima_artifact("precipitation_obs"), "precip.mon.mean.nc"), + "precip", + new_start_date = start_date, + shift_by = Dates.firstdayofmonth, + ) + return obs_var + end + + for (sim_name, obs_name) in sim_obs_short_names_sf + obs_var_dict[sim_name] = + (start_date) -> begin + obs_var = ClimaAnalysis.OutputVar( + joinpath( + @clima_artifact( + "era5_monthly_averages_surface_single_level_1979_2024" + ), + "era5_monthly_averages_surface_single_level_197901-202410.nc", + ), + obs_name, + new_start_date = start_date, + shift_by = Dates.firstdayofmonth, + ) + (ClimaAnalysis.units(obs_var) == "W m**-2") && ( + obs_var = ClimaAnalysis.convert_units( + obs_var, + "W m^-2", + conversion_function = units -> units * -1.0, + ) + ) + return obs_var + end + end + + for (sim_name, obs_name) in sim_obs_short_names_rad + obs_var_dict[sim_name] = + (start_date) -> begin + obs_var = ClimaAnalysis.OutputVar( + joinpath( + @clima_artifact("radiation_obs"), + "CERES_EBAF_Ed4.2_Subset_200003-201910.nc", + ), + obs_name, + new_start_date = start_date, + shift_by = Dates.firstdayofmonth, + ) + # Convert from W m-2 to W m^-2 + ClimaAnalysis.units(obs_var) == "W m-2" ? + obs_var = ClimaAnalysis.set_units(obs_var, "W m^-2") : + error("Did not expect $(ClimaAnalysis.units(obs_var)) for the units") + return obs_var + end + end + return obs_var_dict +end + +""" + get_rmse_var_dict() + +Return a dictionary mapping short names to `RMSEVariable` containing information about +RMSEs of models for the short name of the variable. This is used by the function +`compute_leaderboard`. +""" +function get_rmse_var_dict() + # Load data into RMSEVariables + rmse_var_pr = ClimaAnalysis.read_rmses( + joinpath(@clima_artifact("cmip_model_rmse"), "pr_rmse_amip_pr_amip_5yr.csv"), + "pr", + units = "mm / day", + ) + rmse_var_rsut = ClimaAnalysis.read_rmses( + joinpath(@clima_artifact("cmip_model_rmse"), "rsut_rmse_amip_rsut_amip_5yr.csv"), + "rsut", + units = "W m^-2", + ) + rmse_var_rlut = ClimaAnalysis.read_rmses( + joinpath(@clima_artifact("cmip_model_rmse"), "rlut_rmse_amip_rlut_amip_5yr.csv"), + "rlut", + units = "W m^-2", + ) + + # Add models and units for CliMA + rmse_var_pr = ClimaAnalysis.add_model(rmse_var_pr, "CliMA") + ClimaAnalysis.add_unit!(rmse_var_pr, "CliMA", "mm / day") + + rmse_var_rsut = ClimaAnalysis.add_model(rmse_var_rsut, "CliMA") + ClimaAnalysis.add_unit!(rmse_var_rsut, "CliMA", "W m^-2") + + rmse_var_rlut = ClimaAnalysis.add_model(rmse_var_rlut, "CliMA") + ClimaAnalysis.add_unit!(rmse_var_rlut, "CliMA", "W m^-2") + + # Store the rmse vars in a dictionary + rmse_var_dict = ClimaAnalysis.Var.OrderedDict( + "pr" => rmse_var_pr, + "rsut" => rmse_var_rsut, + "rlut" => rmse_var_rlut, + ) + return rmse_var_dict +end + +""" + get_sim_var_in_pfull_dict(diagnostics_folder_path) + +Return a dictionary mapping short names to `OutputVar` containing preprocessed +simulation data in pressure space. This is used by `compute_pfull_leaderboard`. + +To add a variable for the leaderboard, add a key-value pair to the dictionary +`sim_var_dict` whose key is the short name of the variable and the value is an +anonymous function that returns a `OutputVar`. For each variable, any +preprocessing should be done in the corresponding anonymous function which +includes unit conversion and shifting the dates. + +The variable should have only four dimensions: latitude, longitude, time, and +pressure. The units of pressure should be in hPa. + +Some configurations have pressure inversions at low elevations, which prevents the +conversion to pressure coordiantes. This function discards everything below 80 meters of +elevation to avoid that. +""" +function get_sim_var_in_pfull_dict(diagnostics_folder_path) + available_short_names = + ClimaAnalysis.available_vars(ClimaAnalysis.SimDir(diagnostics_folder_path)) + sim_var_pfull_dict = Dict{String, Any}() + + short_names = get_short_names_monthly_averages(diagnostics_folder_path) + available_short_names = intersect(short_names, Set(["ta", "hur", "hus"])) + for short_name in short_names + short_name in available_short_names && ( + sim_var_pfull_dict[short_name] = + () -> begin + sim_var = get( + ClimaAnalysis.SimDir(diagnostics_folder_path), + short_name = short_name, + reduction = "average", + period = "1M", + ) + pfull_var = get( + ClimaAnalysis.SimDir(diagnostics_folder_path), + short_name = "pfull", + reduction = "average", + period = "1M", + ) + + (ClimaAnalysis.units(sim_var) == "") && + (sim_var = ClimaAnalysis.set_units(sim_var, "unitless")) + (ClimaAnalysis.units(sim_var) == "kg kg^-1") && + (sim_var = ClimaAnalysis.set_units(sim_var, "unitless")) + + # For certain grid configurations, certain columns can have pressure + # inversions at low elevation, which prevents the conversion between + # pressure and altitude. To avoid that, we exclude everything below 80m. + # NOTE: This was empirically found. + pfull_var_windowed = + ClimaAnalysis.window(pfull_var, "z", left = 80) + sim_var_windowed = ClimaAnalysis.window(sim_var, "z", left = 80) + + sim_in_pfull_var = ClimaAnalysis.Atmos.to_pressure_coordinates( + sim_var_windowed, + pfull_var_windowed, + ) + sim_in_pfull_var = + ClimaAnalysis.shift_to_start_of_previous_month(sim_in_pfull_var) + sim_in_pfull_var = ClimaAnalysis.convert_dim_units( + sim_in_pfull_var, + "pfull", + "hPa"; + conversion_function = x -> 0.01 * x, + ) + return sim_in_pfull_var + end + ) + end + return sim_var_pfull_dict +end + +""" + get_obs_var_in_pfull_dict() + +Return a dictionary mapping short names to `OutputVar` containing preprocessed +observational data in pressure coordinates. This is used by +`compute_pfull_leaderboard`. + +To add a variable for the leaderboard, add a key-value pair to the dictionary +`obs_var_dict` whose key is the short name of the variable and the value is an anonymous +function that returns a `OutputVar`. The function must take in a start date +which is used to align the times in the observational data to match the +simulation data. The short name must be the same as in `sim_var_pfull_dict` in +the function `get_sim_var_in_pfull_dict`. Any preprocessing is done in the +function which includes unit conversion and shifting the dates. + +The variable should have only four dimensions: latitude, longitude, time, and +pressure. The units of pressure should be in hPa. +""" +function get_obs_var_in_pfull_dict() + artifact_path = joinpath( + @clima_artifact("era5_monthly_averages_pressure_levels_1979_2024"), + "era5_monthly_averages_pressure_levels_197901-202410.nc", + ) + short_names_pairs = [("ta", "t"), ("hus", "q")] + obs_var_dict = Dict{String, Any}() + for (short_name, era5_short_name) in short_names_pairs + obs_var_dict[short_name] = + (start_date) -> begin + obs_var = ClimaAnalysis.OutputVar( + artifact_path, + era5_short_name, + new_start_date = start_date, + shift_by = Dates.firstdayofmonth, + ) + (ClimaAnalysis.units(obs_var) == "kg kg**-1") && + (obs_var = ClimaAnalysis.set_units(obs_var, "unitless")) + obs_var = ClimaAnalysis.Var.convert_dim_units( + obs_var, + "pressure_level", + "hPa"; + conversion_function = x -> 0.01 * x, + ) + return obs_var + end + end + + obs_var_dict["hur"] = + (start_date) -> begin + obs_var = ClimaAnalysis.OutputVar( + artifact_path, + "r", + new_start_date = start_date, + shift_by = Dates.firstdayofmonth, + ) + obs_var = ClimaAnalysis.Var.convert_dim_units( + obs_var, + "pressure_level", + "hPa"; + conversion_function = x -> 0.01 * x, + ) + # Convert from percentages (e.g. 120%) to decimal (1.20) + obs_var = ClimaAnalysis.Var.convert_units( + obs_var, + "unitless", + conversion_function = x -> 0.01 * x, + ) + return obs_var + end + return obs_var_dict +end + +""" + get_compare_vars_biases_plot_extrema_pfull() + +Return a dictionary mapping short names to ranges for the bias plots. This is +used by the function `compute_pfull_leaderboard`. + +To add a variable to the leaderboard, add a key-value pair to the dictionary +`compare_vars_biases_plot_extrema` whose key is a short name key is the same +short name in `sim_var_pfull_dict` in the function `get_sim_var_pfull_dict` and +the value is a tuple, where the first element is the lower bound and the last +element is the upper bound for the bias plots. +""" +function get_compare_vars_biases_plot_extrema_pfull() + compare_vars_biases_plot_extrema = + Dict("ta" => (-5.0, 5.0), "hur" => (-0.4, 0.4), "hus" => (-0.0015, 0.0015)) + return compare_vars_biases_plot_extrema +end + +""" + get_compare_vars_biases_heatmap_extrema_pfull() + +Return a dictionary mapping short names to ranges for the heatmaps. This is +used by the function `compute_pfull_leaderboard`. + +To add a variable to the leaderboard, add a key-value pair to the dictionary +`compare_vars_biases_heatmap_extrema` whose key is a short name key is the same +short name in `sim_var_pfull_dict` in the function `get_sim_var_pfull_dict` and +the value is a tuple, where the first element is the lower bound and the last +element is the upper bound for the bias plots. +""" +function get_compare_vars_biases_heatmap_extrema_pfull() + compare_vars_biases_heatmap_extrema = + Dict("ta" => (-10.0, 10.0), "hur" => (-0.4, 0.4), "hus" => (-0.001, 0.001)) + return compare_vars_biases_heatmap_extrema +end + +""" + get_short_names_of_monthly_averages(diagnostics_folder_path) + +Get all the short names of the monthly averages. +""" +function get_short_names_monthly_averages(diagnostics_folder_path) + available_short_names = Set{String}() + simdir = ClimaAnalysis.SimDir(diagnostics_folder_path) + for short_name in ClimaAnalysis.available_vars(simdir) + for reduction in ClimaAnalysis.available_reductions(simdir, short_name = short_name) + for period in ClimaAnalysis.available_periods( + simdir, + short_name = short_name, + reduction = reduction, + ) + if reduction == "average" && period == "1M" + push!(available_short_names, short_name) + end + end + end + end + return available_short_names +end diff --git a/src/Postprocessor/leaderboard/leaderboard.jl b/src/Postprocessor/leaderboard/leaderboard.jl new file mode 100644 index 0000000000..04afd40c83 --- /dev/null +++ b/src/Postprocessor/leaderboard/leaderboard.jl @@ -0,0 +1,383 @@ +import ClimaAnalysis +import GeoMakie +import CairoMakie +import Dates + +include("data_sources.jl") + +""" + compute_leaderboard(leaderboard_base_path, diagnostics_folder_path, spinup) + +Plot the biases and a leaderboard of various variables defined over longitude, latitude, and +time. + +The argument `leaderboard_base_path` is the path to save the leaderboards and bias plots, +`diagnostics_folder_path` is the path to the simulation data, and `spinup` is the number +of months to remove from the beginning of the simulation. + +Loading and preprocessing simulation data is done by `get_sim_var_dict`. Loading and +preprocessing observational data is done by `get_obs_var_dict`. The ranges of the bias plots +are determined by `get_compare_vars_biases_plot_extrema`. The groups of variables plotted on +the bias plots are determined by `get_compare_vars_biases_groups()`. Loading the RMSEs from +other models is done by `get_rmse_var_dict`. See the functions defined in data_sources.jl. +""" +function compute_leaderboard(leaderboard_base_path, diagnostics_folder_path, spinup) + @info "Error against observations" + + # Get everything we need from data_sources.jl + sim_var_dict = get_sim_var_dict(diagnostics_folder_path) + obs_var_dict = get_obs_var_dict() + compare_vars_biases_plot_extrema = get_compare_vars_biases_plot_extrema() + rmse_var_dict = get_rmse_var_dict() + compare_vars_biases_groups = get_compare_vars_biases_groups() + + # Set up dict for storing simulation and observational data after processing + sim_obs_comparsion_dict = Dict() + seasons = ["ANN", "MAM", "JJA", "SON", "DJF"] + + # Print dates for debugging + _, var_func = first(sim_var_dict) + var = var_func() + @info "Working with dates:" + @info ClimaAnalysis.dates(var) + + for short_name in keys(sim_var_dict) + # Simulation data + sim_var = sim_var_dict[short_name]() + + # Observational data + obs_var = obs_var_dict[short_name](sim_var.attributes["start_date"]) + + # Remove first spin_up_months from simulation + spinup_cutoff = spinup * 31 * 86400.0 + ClimaAnalysis.times(sim_var)[end] >= spinup_cutoff && + (sim_var = ClimaAnalysis.window(sim_var, "time", left = spinup_cutoff)) + + obs_var = ClimaAnalysis.resampled_as(obs_var, sim_var) + obs_var_seasons = ClimaAnalysis.split_by_season(obs_var) + sim_var_seasons = ClimaAnalysis.split_by_season(sim_var) + + # Add annual to start of seasons + obs_var_seasons = (obs_var, obs_var_seasons...) + sim_var_seasons = (sim_var, sim_var_seasons...) + + # Take time average + obs_var_seasons = ( + !ClimaAnalysis.isempty(obs_var) ? ClimaAnalysis.average_time(obs_var) : obs_var for obs_var in obs_var_seasons + ) + sim_var_seasons = ( + !ClimaAnalysis.isempty(sim_var) ? ClimaAnalysis.average_time(sim_var) : sim_var for sim_var in sim_var_seasons + ) + + # Save observation and simulation data + sim_obs_comparsion_dict[short_name] = Dict( + season => (sim_var_s, obs_var_s) for (season, sim_var_s, obs_var_s) in + zip(seasons, sim_var_seasons, obs_var_seasons) + ) + end + + # Filter seasons to remove seasons with no dates + _, var = first(sim_obs_comparsion_dict) + filter!(season -> !ClimaAnalysis.isempty(var[season][1]), seasons) + + # Plot annual plots + for compare_vars_biases in compare_vars_biases_groups + # Plot all the variables that we can + compare_vars_biases = + filter(x -> x in keys(sim_obs_comparsion_dict), compare_vars_biases) + length(compare_vars_biases) > 0 || continue + + fig_bias = + CairoMakie.Figure(; size = (600, 75.0 + 300 * length(compare_vars_biases))) + for (row, short_name) in enumerate(compare_vars_biases) + sim_var = sim_obs_comparsion_dict[short_name]["ANN"][1] + if !ClimaAnalysis.isempty(sim_var) + # Add "mean " for plotting the title + sim_var.attributes["short_name"] = "mean $(ClimaAnalysis.short_name(sim_var))" + cmap_extrema = get(compare_vars_biases_plot_extrema, short_name) do + extrema( + ClimaAnalysis.bias( + sim_obs_comparsion_dict[short_name]["ANN"]..., + ).data, + ) + end + ClimaAnalysis.Visualize.plot_bias_on_globe!( + fig_bias, + sim_obs_comparsion_dict[short_name]["ANN"]..., + cmap_extrema = cmap_extrema, + p_loc = (row, 1), + ) + end + end + CairoMakie.save( + joinpath(leaderboard_base_path, "bias_$(first(compare_vars_biases))_ANN.png"), + fig_bias, + ) + end + + # Plot plots with all the seasons (not annual) + seasons_no_ann = filter(season -> season != "ANN", seasons) + for compare_vars_biases in compare_vars_biases_groups + # Plot all the variables that we can + compare_vars_biases = + filter(x -> x in keys(sim_obs_comparsion_dict), compare_vars_biases) + length(compare_vars_biases) > 0 || continue + + fig_all_seasons = CairoMakie.Figure(; + size = (600 * length(seasons_no_ann), 75.0 + 300 * length(compare_vars_biases)), + ) + for (col, season) in enumerate(seasons_no_ann) + CairoMakie.Label( + fig_all_seasons[0, col], + season, + tellwidth = false, + fontsize = 30, + ) + for (row, short_name) in enumerate(compare_vars_biases) + sim_var = sim_obs_comparsion_dict[short_name][season][1] + if !ClimaAnalysis.isempty(sim_var) + cmap_extrema = get(compare_vars_biases_plot_extrema, short_name) do + extrema( + ClimaAnalysis.bias( + sim_obs_comparsion_dict[short_name][season]..., + ).data, + ) + end + # Add "mean " for plotting the title + sim_var.attributes["short_name"] = "mean $(ClimaAnalysis.short_name(sim_var))" + ClimaAnalysis.Visualize.plot_bias_on_globe!( + fig_all_seasons[row, col], + sim_obs_comparsion_dict[short_name][season]..., + cmap_extrema = cmap_extrema, + ) + end + end + end + CairoMakie.save( + joinpath( + leaderboard_base_path, + "bias_$(first(compare_vars_biases))_all_seasons.png", + ), + fig_all_seasons, + ) + end + + # Add RMSE for the CliMA model and for each season + is_in_sim_vars(k) = k in keys(sim_var_dict) + rmse_var_dict = Dict(k => v for (k, v) in rmse_var_dict if is_in_sim_vars(k)) + rmse_var_names = keys(rmse_var_dict) + for short_name in rmse_var_names + for season in seasons + !ClimaAnalysis.isempty(sim_obs_comparsion_dict[short_name][season][1]) && ( + rmse_var_dict[short_name]["CliMA", season] = ClimaAnalysis.global_rmse( + sim_obs_comparsion_dict[short_name][season]..., + ) + ) + end + end + + # Plot box plots + fig_leaderboard = + CairoMakie.Figure(; size = (800, 300 * length(rmse_var_dict) + 400), fontsize = 20) + for (loc, short_name) in enumerate(rmse_var_names) + ClimaAnalysis.Visualize.plot_boxplot!( + fig_leaderboard, + rmse_var_dict[short_name], + ploc = (loc, 1), + best_and_worst_category_name = "ANN", + ) + end + + # Plot leaderboard + ClimaAnalysis.Visualize.plot_leaderboard!( + fig_leaderboard, + (rmse_var_dict[short_name] for short_name in rmse_var_names)..., + best_category_name = "ANN", + ploc = (length(rmse_var_dict) + 1, 1), + ) + CairoMakie.save( + joinpath(leaderboard_base_path, "bias_leaderboard.png"), + fig_leaderboard, + ) +end + +""" + compute_pfull_leaderboard(leaderboard_base_path, diagnostics_folder_path, spinup) + +Plot the biases and a leaderboard of various variables defined over longitude, latitude, +time, and pressure levels. + +The argument `leaderboard_base_path` is the path to save the leaderboards and bias plots, +`diagnostics_folder_path` is the path to the simulation data, and `spinup` is the number +of months to remove from the beginning of the simulation. + +Loading and preprocessing simulation data is done by `get_sim_var_in_pfull_dict`. Loading +and preprocessing observational data is done by `get_obs_var_in_pfull_dict`. The ranges of +the bias plots is defined by `get_compare_vars_biases_plot_extrema_pfull`. See the functions +defined in data_sources.jl for more information. +""" +function compute_pfull_leaderboard(leaderboard_base_path, diagnostics_folder_path, spinup) + @info "Error against observations for variables in pressure coordinates" + + sim_var_pfull_dict = get_sim_var_in_pfull_dict(diagnostics_folder_path) + obs_var_pfull_dict = get_obs_var_in_pfull_dict() + + # Print dates for debugging + _, get_var = first(sim_var_pfull_dict) + var = get_var() + @info "Working with dates:" + @info ClimaAnalysis.dates(var) + + # Set up dict for storing simulation and observational data after processing + sim_obs_comparsion_dict = Dict() + + for short_name in keys(sim_var_pfull_dict) + sim_var = sim_var_pfull_dict[short_name]() + obs_var = obs_var_pfull_dict[short_name](sim_var.attributes["start_date"]) + + # Check units for pressure are in hPa + ClimaAnalysis.dim_units(sim_var, ClimaAnalysis.pressure_name(sim_var)) != "hPa" && + error("Units of pressure should be hPa for $short_name simulation data") + ClimaAnalysis.dim_units(obs_var, ClimaAnalysis.pressure_name(obs_var)) != "hPa" && + error("Units of pressure should be hPa for $short_name simulation data") + + # Remove first spin_up_months from simulation + spinup_cutoff = spinup * 31 * 86400.0 + ClimaAnalysis.times(sim_var)[end] >= spinup_cutoff && + (sim_var = ClimaAnalysis.window(sim_var, "time", left = spinup_cutoff)) + + # Restrain the pressure levels so we can resample + min_pfull = ClimaAnalysis.pressures(obs_var)[1] + sim_pressures = ClimaAnalysis.pressures(sim_var) + greater_than_min_pfull_idx = findfirst(x -> x > min_pfull, sim_pressures) + sim_var = ClimaAnalysis.window( + sim_var, + "pfull", + left = sim_pressures[greater_than_min_pfull_idx], + ) + + # Resample over lat, lon, time, and pressures + obs_var = ClimaAnalysis.resampled_as(obs_var, sim_var) + + # Take time average + sim_var = ClimaAnalysis.average_time(sim_var) + obs_var = ClimaAnalysis.average_time(obs_var) + + # Save observation and simulation data + sim_obs_comparsion_dict[short_name] = (; sim = sim_var, obs = obs_var) + end + + # Slicing only uses the nearest value, so we also need to keep track of the + # actual pressure levels that we get when slicing + target_p_lvls = [850.0, 500.0, 250.0] + real_p_lvls = [] + + # Get units for pressure for plotting + p_units = ClimaAnalysis.dim_units(first(sim_obs_comparsion_dict)[2].sim, "pfull") + + # Initialize ranges for colorbar and figure whose columns are pressure levels and rows are variables + compare_vars_biases_plot_extrema_pfull = get_compare_vars_biases_plot_extrema_pfull() + fig_bias_pfull_vars = CairoMakie.Figure( + size = (650 * length(target_p_lvls), 450 * length(sim_obs_comparsion_dict)), + ) + + # Plot bias at the pressure levels in p_lvls + for (row_idx, short_name) in enumerate(keys(sim_obs_comparsion_dict)) + # Plot label for variable name + CairoMakie.Label( + fig_bias_pfull_vars[row_idx, 0], + short_name, + tellheight = false, + fontsize = 30, + ) + for (col_idx, p_lvl) in enumerate(target_p_lvls) + sim_var = sim_obs_comparsion_dict[short_name].sim + obs_var = sim_obs_comparsion_dict[short_name].obs + + # Slice at pressure level using nearest value rather interpolating + sim_var = ClimaAnalysis.slice(sim_var, pfull = p_lvl) + obs_var = ClimaAnalysis.slice(obs_var, pressure_level = p_lvl) + + # Get the actual pressure level that we slice at + push!(real_p_lvls, parse(Float64, sim_var.attributes["slice_pfull"])) + + # Add so that it show up on in the title + sim_var.attributes["short_name"] = "mean $(ClimaAnalysis.short_name(sim_var))" + + # Plot bias + ClimaAnalysis.Visualize.plot_bias_on_globe!( + fig_bias_pfull_vars[row_idx, col_idx], + sim_var, + obs_var, + cmap_extrema = compare_vars_biases_plot_extrema_pfull[short_name], + ) + end + end + + # Plot label for the pressure levels + for (col_idx, p_lvl) in enumerate(real_p_lvls[begin:length(target_p_lvls)]) + CairoMakie.Label( + fig_bias_pfull_vars[0, col_idx], + "$p_lvl $p_units", + tellwidth = false, + fontsize = 30, + ) + end + + # Define figure with at most four columns and the necessary number of rows for + # all the variables + num_vars = length(sim_obs_comparsion_dict) + num_cols = min(4, num_vars) + num_rows = ceil(Int, num_vars / 4) + fig_lat_pres = CairoMakie.Figure(size = (650 * num_cols, 450 * num_rows)) + + # Initialize ranges for colorbar + compare_vars_biases_heatmap_extrema_pfull = + get_compare_vars_biases_heatmap_extrema_pfull() + + # Take zonal mean and plot lat - pressure heatmap + for (idx, short_name) in enumerate(keys(sim_obs_comparsion_dict)) + # Compute where the figure should be placed + # Goes from 1 -> (1,1), 2 -> (1,2), ..., 4 -> (1,4), 5 -> (2,1) + # for idx -> (col_idx, row_idx) + row_idx = div(idx - 1, 4) + 1 + col_idx = 1 + rem(idx - 1, 4) + + # Load data + sim_var = sim_obs_comparsion_dict[short_name].sim + obs_var = sim_obs_comparsion_dict[short_name].obs + + # Take zonal mean (average along lon arithmetically) + sim_var = ClimaAnalysis.average_lon(sim_var) + obs_var = ClimaAnalysis.average_lon(obs_var) + + # Compute bias and set short name, long name, and units + bias_var = sim_var - obs_var + bias_var = ClimaAnalysis.set_units(bias_var, ClimaAnalysis.units(sim_var)) + bias_var.attributes["short_name"] = "sim-obs_$(ClimaAnalysis.short_name(sim_var))" + bias_var.attributes["long_name"] = "SIM-OBS_$(ClimaAnalysis.long_name(sim_var))" + ClimaAnalysis.Visualize.heatmap2D!( + fig_lat_pres[row_idx, col_idx], + bias_var, + more_kwargs = Dict( + :axis => Dict([:yreversed => true]), + :plot => Dict( + :colormap => :vik, + :colorrange => + compare_vars_biases_heatmap_extrema_pfull[short_name], + :colormap => CairoMakie.cgrad(:vik, 11, categorical = true), + ), + ), + ) + end + + # Save figures + CairoMakie.save( + joinpath(leaderboard_base_path, "bias_vars_in_pfull.png"), + fig_bias_pfull_vars, + ) + CairoMakie.save( + joinpath(leaderboard_base_path, "bias_lat_pfull_heatmaps.png"), + fig_lat_pres, + ) +end diff --git a/src/Postprocessor/leaderboard/test_rmses.jl b/src/Postprocessor/leaderboard/test_rmses.jl new file mode 100644 index 0000000000..9d9e39c0e5 --- /dev/null +++ b/src/Postprocessor/leaderboard/test_rmses.jl @@ -0,0 +1,70 @@ +import ClimaAnalysis +import Dates +import Test: @test, @testset + +include("data_sources.jl") + +""" + test_rmse_thresholds(diagnostics_folder_path, spinup) + +Test that the annual RMSE values for specific variables have not increased +beyond acceptable thresholds. The variables tested are: +- pr (precipitation) +- rsut (top of atmosphere outgoing shortwave radiation) +- rsutcs (clear-sky top of atmosphere outgoing shortwave radiation) + +The spinup is the number of months to remove from the beginning of the +simulation. + +More variables can be added by adding the short name and RMSE pair to the +dictionary returned by `get_rmse_thresholds`. + +If this test fails, it indicates a regression in the model's physics, resulting +in a higher RMSE. If this increased RMSE is considered acceptable, then the +thresholds should be updated accordingly. +""" +function test_rmse_thresholds(diagnostics_folder_path, spinup) + sim_var_dict = get_sim_var_dict(diagnostics_folder_path) + obs_var_dict = get_obs_var_dict() + rmse_thresholds = get_rmse_thresholds() + + sim_vars = (sim_var_dict[short_name]() for short_name in keys(rmse_thresholds)) + obs_vars = ( + obs_var_dict[ClimaAnalysis.short_name(sim_var)](sim_var.attributes["start_date"]) for sim_var in sim_vars + ) + short_names = (ClimaAnalysis.short_name(var) for var in sim_vars) + + rmses = map(sim_vars, obs_vars) do sim_var, obs_var + # Remove first spin_up_months from simulation + spinup_cutoff = spinup * 31 * 86400.0 + ClimaAnalysis.times(sim_var)[end] >= spinup_cutoff && + (sim_var = ClimaAnalysis.window(sim_var, "time", left = spinup_cutoff)) + + obs_var = ClimaAnalysis.resampled_as(obs_var, sim_var) + obs_var = ClimaAnalysis.average_time(obs_var) + sim_var = ClimaAnalysis.average_time(sim_var) + + ClimaAnalysis.global_rmse(sim_var, obs_var) + end + + @testset "RMSE thresholds" begin + for (short_name, rmse) in zip(short_names, rmses) + @info "RMSE for $short_name: $rmse" + @test rmse < rmse_thresholds[short_name] + end + end +end + +""" + get_rmse_thresholds() + +Return a dictionary mapping short names to maximum acceptable RMSE values. +""" +function get_rmse_thresholds() + rmse_thresholds = Dict( + "pr" => 3.0, # mm/day + "rsut" => 29.0, # W/m² + "rsutcs" => 10.8, # W/m² + ) + return rmse_thresholds +end diff --git a/src/SimCoordinator.jl b/src/SimCoordinator.jl new file mode 100644 index 0000000000..f31d0c31ad --- /dev/null +++ b/src/SimCoordinator.jl @@ -0,0 +1,114 @@ +""" + SimCoordinator + +This module handles timestepping and running of coupled simulations. +It coordinates the execution of component models and manages the simulation lifecycle. +""" +module SimCoordinator + +import ClimaComms +import ..Interfacer +import ..ConservationChecker +import ..FieldExchanger +import ..FluxCalculator +import ..TimeManager +import ..Postprocessor + +export run!, step! + +""" + run!(cs::Interfacer.CoupledSimulation; precompile = ...) + +Run the coupled simulation for the full timespan. + +Keyword arguments +================== + +`precompile`: If `true`, run the coupled simulations for two steps, so that most functions + are precompiled and subsequent timing will be more accurate. +""" +function run!( + cs::Interfacer.CoupledSimulation; + precompile = (cs.tspan[end] > 2 * cs.Δt_cpl + cs.tspan[begin]), +) + ## Precompilation of Coupling Loop + # Here we run the entire coupled simulation for two timesteps to precompile several + # functions for more accurate timing of the overall simulation. + precompile && (step!(cs); step!(cs)) + + ## Run garbage collection before solving for more accurate memory comparison to ClimaAtmos + GC.gc() + + #= + ## Solving and Timing the Full Simulation + + This is where the full coupling loop, `solve_coupler!` is called for the full timespan of the simulation. + We use the `ClimaComms.@elapsed` macro to time the simulation on both CPU and GPU, and use this + value to calculate the simulated years per day (SYPD) of the simulation. + =# + @info "Starting coupling loop" + # Note: ClimaAtmos timing and helper functions are expected to be available in the calling context + # This allows SimCoordinator to remain independent while still supporting ClimaEarth-specific functionality + walltime = ClimaComms.@elapsed ClimaComms.device(cs) begin + while cs.t[] < cs.tspan[end] + step!(cs) + end + end + @info "Simulation took $(walltime) seconds" + + # Compute and log performance metrics + sypd = Postprocessor.simulated_years_per_day(cs, walltime) + walltime_per_step = Postprocessor.walltime_per_coupling_step(cs, walltime) + @info "SYPD: $sypd" + @info "Walltime per coupling step: $(walltime_per_step)" + Postprocessor.save_sypd_walltime_to_disk(cs, walltime) + + # Close all diagnostics file writers + isnothing(cs.diags_handler) || + foreach(diag -> close(diag.output_writer), cs.diags_handler.scheduled_diagnostics) + foreach(Interfacer.close_output_writers, cs.model_sims) + + return nothing +end + +""" + step!(cs::Interfacer.CoupledSimulation) + +Take one coupling step forward in time. + +This function runs the component models sequentially, and exchanges combined fields and +calculates fluxes using the selected turbulent fluxes option. Note, one coupling step might +require multiple steps in some of the component models. +""" +function step!(cs::Interfacer.CoupledSimulation) + # Update the current time + cs.t[] += cs.Δt_cpl + + ## compute global energy and water conservation checks + ## (only for slabplanet if tracking conservation is enabled) + ConservationChecker.check_conservation!(cs) + + ## step component model simulations sequentially for one coupling timestep (Δt_cpl) + FieldExchanger.step_model_sims!(cs) + + ## update the surface fractions for surface models + FieldExchanger.update_surface_fractions!(cs) + + ## exchange all non-turbulent flux fields between models, including radiative and precipitation fluxes + FieldExchanger.exchange!(cs) + + ## calculate turbulent fluxes in the coupler and update the model simulations with them + FluxCalculator.turbulent_fluxes!(cs) + + ## compute any ocean-sea ice fluxes + FluxCalculator.ocean_seaice_fluxes!(cs) + + ## Maybe call the callbacks + TimeManager.callbacks!(cs) + + # Compute coupler diagnostics + Postprocessor.CD.orchestrate_diagnostics(cs) + return nothing +end + +end # module