use ClimaDiagnostics for AMIP diags

experiments/ClimaEarth/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.5"
 manifest_format = "2.0"
-project_hash = "0a5ef4691aac8519cb45950f67c96bcaab893b0a"
+project_hash = "4358de6a1c73febd571625d02f1a6682a5d50335"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "99a6f5d0ce1c7c6afdb759daa30226f71c54f6b0"

experiments/ClimaEarth/Project.toml

@@ -9,6 +9,7 @@ ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaCorePlots = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 ClimaCoupler = "4ade58fe-a8da-486c-bd89-46df092ec0c7"
+ClimaDiagnostics = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
 ClimaLand = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"

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

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