diff --git a/experiments/ClimaEarth/run_amip.jl b/experiments/ClimaEarth/run_amip.jl index 30451e5501..a58cab4b6a 100644 --- a/experiments/ClimaEarth/run_amip.jl +++ b/experiments/ClimaEarth/run_amip.jl @@ -48,14 +48,7 @@ import ClimaCore as CC # ## Coupler specific imports import ClimaCoupler import ClimaCoupler: - ConservationChecker, - Checkpointer, - FieldExchanger, - FluxCalculator, - Interfacer, - Regridder, - TimeManager, - Utilities + ConservationChecker, Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Regridder, TimeManager, Utilities import ClimaUtilities.SpaceVaryingInputs: SpaceVaryingInput import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput, evaluate! @@ -81,6 +74,7 @@ include("components/ocean/eisenman_seaice.jl") include("user_io/user_logging.jl") include("user_io/debug_plots.jl") include("user_io/io_helpers.jl") +mode_name == "amip" && include("user_io/amip_diagnostics.jl") #= ### Configuration Dictionaries @@ -561,6 +555,11 @@ cs = Interfacer.CoupledSimulation{FT}( ); Utilities.show_memory_usage(comms_ctx) +#= Set up default AMIP diagnostics +Use ClimaDiagnostics for default AMIP diagnostics, which currently include turbulent energy fluxes. +=# +mode_name == "amip" && amip_diags_handler, cs_nt = amip_diagnostics_setup(cs) + #= ## Restart component model states if specified If a restart directory is specified and contains output files from the `checkpoint_cb` callback, the component model states are restarted from those files. The restart directory @@ -712,6 +711,9 @@ function solve_coupler!(cs) ## callback to checkpoint model state TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint) + + ## compute/output AMIP diagnostics if scheduled for this timestep + CD.orchestrate_diagnostics(cs_nt, amip_diags_handler) end return nothing end @@ -828,7 +830,6 @@ if ClimaComms.iamroot(comms_ctx) ## ClimaESM # Include the user-defined plotting functions and diagnostics include("user_io/ci_plots.jl") - include("user_io/user_diagnostics.jl") amip_short_names = ["ta", "ua", "hus", "clw", "pr", "ts", "toa_fluxes_net", "F_turb_energy"] make_ci_plots([atmos_sim.integrator.p.output_dir], dir_paths.artifacts, short_names = amip_short_names) @@ -957,4 +958,6 @@ if ClimaComms.iamroot(comms_ctx) rm(dir_paths.output; recursive = true, force = true) #hide end #hide + ## close all diagnostics file writers + map(diag -> close(diag.output_writer), amip_diags_handler.scheduled_diagnostics) end diff --git a/experiments/ClimaEarth/user_io/amip_diagnostics.jl b/experiments/ClimaEarth/user_io/amip_diagnostics.jl new file mode 100644 index 0000000000..00ead3541a --- /dev/null +++ b/experiments/ClimaEarth/user_io/amip_diagnostics.jl @@ -0,0 +1,67 @@ +import ClimaDiagnostics as CD +import ClimaCoupler: Interfacer +import Dates + +""" + amip_diagnostics(cs::Interfacer.CoupledSimulation) + +Create and return a list of default diagnostics for AMIP simulations, +using ClimaDiagnostics The diagnostics are saved to NetCDF files. +Currently, this just includes a diagnostic for turbulent energy fluxes. +""" +function amip_diagnostics(cs::Interfacer.CoupledSimulation) + # Create schedules and writer + schedule_everystep = CD.Schedules.EveryStepSchedule(Dates.Hour(1)) + schedule_12hour = CD.Schedules.EveryCalendarDtSchedule(Dates.Hour(12)) + netcdf_writer = CD.Writers.NetCDFWriter(axes(cs.fields.F_turb_energy), output_dir = cs.dirs.artifacts_dir) + average = (x) -> sum(x) / length(x) + + # 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 = "When using partitioned surface fluxes, turbulent energy fluxes are calculated as + the sum of sensible and latent heat fluxes, weighted by surface simulation area. + When using combined surface fluxes, turbulent energy fluxes are calculated using + the surface conditions calculated in ClimaAtmos.", + compute! = (out, state, cache, time) -> begin + if isnothing(out) + return state.fields.F_turb_energy + else + out .= state.fields.F_turb_energy + end + end, + ) + + # Schedule the turbulent energy fluxes to save at every step, and output the average every 12 hours + F_turb_energy_diag_sched = CD.ScheduledDiagnostic( + F_turb_energy_diag, + compute_schedule_func = schedule_everystep, + output_schedule_func = schedule_12hour, + output_writer = netcdf_writer, + reduction_time_func = average, + descriptive_short_name = "F_turb_energy_12hourly", + descriptive_long_name = "Turbulent energy fluxes calculated in the coupler, averaged every 12 hours of simulation", + ) + + return [F_turb_energy_diag_sched] +end + +""" + amip_diagnostics_setup(cs::Interfacer.CoupledSimulation) + +Set up the default diagnostics for an AMIP simulation. Return a DiagnosticsHandler +object to coordinate the diagnostics, as well as a NamedTuple containing simulation +information strucuted as expected by ClimaDiagnostics. +""" +function amip_diagnostics_setup(cs::Interfacer.CoupledSimulation) + scheduled_diags = amip_diagnostics(cs) + amip_diags_handler = CD.DiagnosticsHandler(scheduled_diags, cs, nothing, cs.t) + + # Wrap the CoupledSimulation object in a NamedTuple to match the ClimaDiagnostics interface + cs_nt = (; out = nothing, u = cs, p = nothing, t = cs.t) + + return amip_diags_handler, cs_nt +end diff --git a/experiments/ClimaEarth/user_io/user_diagnostics.jl b/experiments/ClimaEarth/user_io/user_diagnostics.jl deleted file mode 100644 index 8e07b1325a..0000000000 --- a/experiments/ClimaEarth/user_io/user_diagnostics.jl +++ /dev/null @@ -1,28 +0,0 @@ -import ClimaCore as CC -import ClimaAtmos.Parameters as CAP -import ClimaDiagnostics as CD -import Thermodynamics as TD -import ClimaCoupler: Interfacer, Utilities - -## Coupler-specific diagnostics -function compute_F_turb_energy!(cs::Interfacer.CoupledSimulation) - cs_nt = (; out = nothing, u = cs, p = nothing, t = cs.t) -end - -CD.add_diagnostic_variable!( - short_name = "F_turb_energy", - long_name = "Turbulent energy fluxes", - standard_name = "F_turb_energy", - units = "W m^-2", - comments = "When using partitioned surface fluxes, turbulent energy fluxes are calculated as - the sum of sensible and latent heat fluxes, weighted by surface simulation area. - When using combined surface fluxes, turbulent energy fluxes are calculated using - the surface conditions calculated in ClimaAtmos.", - compute! = (out, state, cache, time) -> begin - if isnothing(out) - return state.fields.F_turb_energy - else - out .= state.fields.F_turb_energy - end - end, -)