use ClimaDiagnostics for AMIP diags

The ClimaCoupler Diagnostics module had become redundant with ClimaDiagnostics.jl, a package designed to provide robust diagnostics across the CliMA ecosystem. Here we remove ClimaCoupler.Diagnostics and instead use ClimaDiagnostics. We're able to retrive most of the diagnostics we want directly from ClimaAtmos and ClimaLand, but also want some that come from coupler-computed quantities, such as `F_turb_energy`. In this PR we add this coupler quantity to our output diagnostics using the ClimaDiagnostics interface. This PR also removes the AMIP paperplots function, but this functionality is replaced by the generalized `make_plots` function.
CliMA · Oct 15, 2024 · 4373db6 · 4373db6
1 parent 42854a4
commit 4373db6
Show file tree

Hide file tree

Showing 31 changed files with 201 additions and 862 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -13,3 +13,19 @@ low-resolution version of such files is automatically downloaded when a
 higher-resolution version is not available. Please, refer to
 [ClimaArtifacts](https://github.com/CliMA/ClimaArtifacts) for more information.
 
+### Code cleanup
+#### Remove ClimaCoupler.Diagnostics module - PR [#953](https://github.com/CliMA/ClimaCoupler.jl/pull/953)
+The ClimaCoupler Diagnostics module had become redundant with
+ClimaDiagnostics.jl, a package designed to provide robust
+diagnostics across the CliMA ecosystem.
+Here we remove ClimaCoupler.Diagnostics and instead use
+ClimaDiagnostics.
+
+We're able to retrieve most of the diagnostics
+we want directly from ClimaAtmos and ClimaLand, but also want
+some that come from coupler-computed quantities, such as
+`F_turb_energy`. In this PR we add this coupler quantity
+to our output diagnostics using the ClimaDiagnostics interface.
+
+This PR also removes the AMIP paperplots function, but this
+functionality is replaced by the generalized `make_plots` function.
diff --git a/config/ci_configs/slabplanet_atmos_diags.yml b/config/ci_configs/slabplanet_atmos_diags.yml
@@ -9,12 +9,11 @@ h_elem: 4
 mode_name: "slabplanet"
 moist: "equil"
 mono_surface: true
-output_default_diagnostics: false
 precip_model: "0M"
 rad: "gray"
 t_end: "10days"
 vert_diff: "true"
 diagnostics:
-  - short_name: [mse, lr, edt, evu, ts, mass_strf, stab, vt, egr]
+  - short_name: [mse, lr, edt, evu, ts, mass_strf, stab, vt, egr, toa_fluxes_net]
     reduction_time: average
     period: 1days
diff --git a/docs/make.jl b/docs/make.jl
@@ -59,7 +59,6 @@ experiment_pages = [
 interface_pages = [
     "checkpointer.md",
     "conservation.md",
-    "diagnostics.md",
     "fieldexchanger.md",
     "fluxcalculator.md",
     "interfacer.md",

diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md
diff --git a/docs/src/postprocessor.md b/docs/src/postprocessor.md
@@ -1,10 +1,10 @@
 # PostProcessor
 
-This module contains functions for postprocessing model data that was saved during the simulation
-by `ClimaCoupler.Diagnostics`. This module is used for offline regridding, slicing and spatial
+This module contains functions for postprocessing model data that was saved during the simulation.
+This module is used for offline regridding, slicing and spatial
 averages. It can also handle data from other sources (e.g., NCEP reanalysis).
 
-## Diagnostics API
+## PostProcessor API
 
 ```@docs
 ClimaCoupler.PostProcessor.postprocess
@@ -17,4 +17,3 @@ ClimaCoupler.PostProcessor.RawData
 ClimaCoupler.PostProcessor.DataPackage
 
 ```
-
diff --git a/experiments/ClimaEarth/Project.toml b/experiments/ClimaEarth/Project.toml
@@ -37,7 +37,7 @@ ArtifactWrappers = "0.2"
 ClimaAnalysis = "0.5.10"
 ClimaAtmos = "0.27"
 ClimaCorePlots = "0.2"
-ClimaDiagnostics = "0.2"
+ClimaDiagnostics = "0.2.6"
 ClimaLand = "0.14, 0.15"
 ClimaParams = "0.10"
 ClimaTimeSteppers = "0.7"

diff --git a/experiments/ClimaEarth/components/atmosphere/climaatmos.jl b/experiments/ClimaEarth/components/atmosphere/climaatmos.jl
@@ -4,9 +4,11 @@ import Statistics
 import LinearAlgebra
 import ClimaAtmos as CA
 import ClimaAtmos: set_surface_albedo!
+import ClimaAtmos.Parameters as CAP
 import ClimaCore as CC
 import ClimaCore.Geometry: ⊗
 import SurfaceFluxes as SF
+import Thermodynamics as TD
 import ClimaCoupler: Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Utilities
 
 include("climaatmos_extra_diags.jl")

diff --git a/experiments/ClimaEarth/components/atmosphere/climaatmos_extra_diags.jl b/experiments/ClimaEarth/components/atmosphere/climaatmos_extra_diags.jl
@@ -136,6 +136,29 @@ CAD.add_diagnostic_variable!(
     end,
 )
 
+CAD.add_diagnostic_variable!(
+    short_name = "toa_fluxes_net",
+    long_name = "Net TOA radiation fluxes",
+    standard_name = "net_toa_radiation_fluxes",
+    units = "W m^-2",
+    comments = "Net top of the atmosphere radiation fluxes, calculated by summing the upward and downward shortwave and longwave radiation fluxes.",
+    compute! = (out, state, cache, time) -> begin
+        # Perform sum of radiation fluxes (rsu + rlu - rsd - rld)
+        if isnothing(out)
+            return CC.Fields.array2field(cache.radiation.rrtmgp_model.face_sw_flux_up, axes(state.f)) .+
+                   CC.Fields.array2field(cache.radiation.rrtmgp_model.face_lw_flux_up, axes(state.f)) .-
+                   CC.Fields.array2field(cache.radiation.rrtmgp_model.face_sw_flux_dn, axes(state.f)) .-
+                   CC.Fields.array2field(cache.radiation.rrtmgp_model.face_lw_flux_dn, axes(state.f))
+        else
+            out .=
+                CC.Fields.array2field(cache.radiation.rrtmgp_model.face_sw_flux_up, axes(state.f)) .+
+                CC.Fields.array2field(cache.radiation.rrtmgp_model.face_lw_flux_up, axes(state.f)) .-
+                CC.Fields.array2field(cache.radiation.rrtmgp_model.face_sw_flux_dn, axes(state.f)) .-
+                CC.Fields.array2field(cache.radiation.rrtmgp_model.face_lw_flux_dn, axes(state.f))
+        end
+    end,
+)
+
 """
     static_stability(cache)
 

diff --git a/experiments/ClimaEarth/run_amip.jl b/experiments/ClimaEarth/run_amip.jl
@@ -48,15 +48,7 @@ import ClimaCore as CC
 # ## Coupler specific imports
 import ClimaCoupler
 import ClimaCoupler:
-    ConservationChecker,
-    Checkpointer,
-    Diagnostics,
-    FieldExchanger,
-    FluxCalculator,
-    Interfacer,
-    Regridder,
-    TimeManager,
-    Utilities
+    ConservationChecker, Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Regridder, TimeManager, Utilities
 
 import ClimaUtilities.SpaceVaryingInputs: SpaceVaryingInput
 import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput, evaluate!
@@ -70,7 +62,7 @@ pkg_dir = pkgdir(ClimaCoupler)
 
 #=
 ### Helper Functions
-These will be eventually moved to their respective component model and diagnostics packages, and so they should not
+These will be eventually moved to their respective component model and utility packages, and so they should not
 contain any internals of the ClimaCoupler source code, except extensions to the Interfacer functions.
 =#
 
@@ -82,7 +74,6 @@ include("components/ocean/prescr_seaice.jl")
 include("components/ocean/eisenman_seaice.jl")
 
 ## helpers for user-specified IO
-include("user_io/user_diagnostics.jl")
 include("user_io/user_logging.jl")
 include("user_io/debug_plots.jl")
 include("user_io/io_helpers.jl")
@@ -126,6 +117,19 @@ Random.seed!(random_seed)
 atmos_config_dict, config_dict = get_atmos_config_dict(config_dict, job_id)
 atmos_config_object = CA.AtmosConfig(atmos_config_dict)
 
+if config_dict["mode_name"] == "amip" && config_dict["use_coupler_diagnostics"]
+    @info "Using default AMIP diagnostics"
+    !haskey(config_dict, "diagnostics") && (config_dict["diagnostics"] = Vector{Dict{Any, Any}}())
+    push!(
+        config_dict["diagnostics"],
+        Dict(
+            "short_name" => ["mse", "lr", "edt", "evu", "ts", "mass_strf", "stab", "vt", "egr", "toa_fluxes_net"],
+            "reduction_time" => "average",
+            "period" => "1days",
+        ),
+    )
+end
+
 ## read in some parsed command line arguments, required by this script
 mode_name = config_dict["mode_name"]
 energy_check = config_dict["energy_check"]
@@ -145,7 +149,6 @@ restart_dir = config_dict["restart_dir"]
 restart_t = Int(config_dict["restart_t"])
 evolving_ocean = config_dict["evolving_ocean"]
 dt_rad = config_dict["dt_rad"]
-use_coupler_diagnostics = config_dict["use_coupler_diagnostics"]
 use_land_diagnostics = config_dict["use_land_diagnostics"]
 
 #=
@@ -490,37 +493,6 @@ model_sims = (atmos_sim = atmos_sim, ice_sim = ice_sim, land_sim = land_sim, oce
 ## dates
 dates = (; date = [date], date0 = [date0], date1 = [Dates.firstdayofmonth(date0)], new_month = [false])
 
-#=
-### Online Diagnostics
-The user can write custom diagnostics in the `user_diagnostics.jl` file.
-Note, this will be replaced by the diagnostics framework currently in ClimaAtmos, once it is abstracted
-into a more general package, so we can use it to save fields from surface models.
-=#
-if use_coupler_diagnostics
-    monthly_3d_diags = Diagnostics.init_diagnostics(
-        (:T, :u, :q_tot, :q_liq_ice),
-        atmos_sim.domain.center_space;
-        save = TimeManager.Monthly(),
-        operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
-        output_dir = dir_paths.output,
-        name_tag = "monthly_mean_3d_",
-    )
-
-    monthly_2d_diags = Diagnostics.init_diagnostics(
-        (:precipitation_rate, :toa_fluxes, :T_sfc, :turbulent_energy_fluxes),
-        boundary_space;
-        save = TimeManager.Monthly(),
-        operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
-        output_dir = dir_paths.output,
-        name_tag = "monthly_mean_2d_",
-    )
-
-    diagnostics = (monthly_3d_diags, monthly_2d_diags)
-    Utilities.show_memory_usage()
-else
-    diagnostics = ()
-end
-
 #=
 ## Initialize Conservation Checks
 
@@ -594,13 +566,23 @@ else
     error("turb_flux_partition must be either PartitionedStateFluxes or CombinedStateFluxesMOST")
 end
 
+#= Set up default AMIP diagnostics
+Use ClimaDiagnostics for default AMIP diagnostics, which currently include turbulent energy fluxes.
+=#
+if mode_name == "amip"
+    include("user_io/amip_diagnostics.jl")
+    amip_diags_handler = amip_diagnostics_setup(coupler_fields, dir_paths.artifacts, dates.date0[1], tspan[1])
+else
+    amip_diags_handler = nothing
+end
+
 #=
 ## Initialize Coupled Simulation
 
 The coupled simulation is initialized here and saved in a global `CoupledSimulation` struct, `cs`. It contains all the information
 required to run the coupled simulation, including the communication context, the dates, the boundary space, the coupler fields, the
 configuration dictionary, the conservation checks, the time span, the time step, the land fraction, the model simulations, the mode
-specifics, the diagnostics, the callbacks, and the directory paths.
+specifics, the callbacks, the directory paths, and diagnostics for AMIP simulations.
 =#
 
 cs = Interfacer.CoupledSimulation{FT}(
@@ -615,11 +597,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     Δt_cpl,
     model_sims,
     mode_specifics,
-    diagnostics,
     callbacks,
     dir_paths,
     turbulent_fluxes,
     thermo_params,
+    amip_diags_handler,
 );
 Utilities.show_memory_usage()
 
@@ -724,15 +706,6 @@ function solve_coupler!(cs)
             current_CO2 = zeros(boundary_space)
             evaluate!(current_CO2, cs.mode.CO2_timevaryinginput, t)
             Interfacer.update_field!(atmos_sim, Val(:co2), current_CO2)
-
-            ## calculate and accumulate diagnostics at each timestep, if we're using diagnostics in this run
-            if !isempty(cs.diagnostics)
-                ClimaComms.barrier(comms_ctx)
-                Diagnostics.accumulate_diagnostics!(cs)
-
-                ## save and reset monthly averages
-                Diagnostics.save_diagnostics(cs)
-            end
         end
 
         ## compute global energy and water conservation checks
@@ -781,6 +754,11 @@ 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
+        ## wrap the current CoupledSimulation fields and time in a NamedTuple to match the ClimaDiagnostics interface
+        cs_nt = (; u = cs.fields, p = nothing, t = t, step = round(t / Δt_cpl))
+        cs.mode.name == "amip" && CD.orchestrate_diagnostics(cs_nt, cs.amip_diags_handler)
     end
     return nothing
 end
@@ -897,9 +875,25 @@ if ClimaComms.iamroot(comms_ctx)
         ## ClimaESM
         include("user_io/ci_plots.jl")
 
-        # TODO add turbulent_energy_fluxes?
-        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)
+        # define variable names and output directories for each diagnostic
+        amip_short_names_atmos = ["ta", "ua", "hus", "clw", "pr", "ts", "toa_fluxes_net"]
+        output_dir_atmos = atmos_sim.integrator.p.output_dir
+        amip_short_names_coupler = ["F_turb_energy"]
+        output_dir_coupler = dir_paths.output
+
+        # Check if all output variables are available in the specified directories
+        make_ci_plots(
+            output_dir_atmos,
+            dir_paths.artifacts,
+            short_names = amip_short_names_atmos,
+            output_prefix = "atmos_",
+        )
+        make_ci_plots(
+            output_dir_coupler,
+            dir_paths.artifacts,
+            short_names = amip_short_names_coupler,
+            output_prefix = "coupler_",
+        )
 
         # Check this because we only want monthly data for making plots
         if t_end > 84600 * 31 * 3 && config_dict["output_default_diagnostics"]
@@ -913,7 +907,7 @@ if ClimaComms.iamroot(comms_ctx)
     if config_dict["ci_plots"]
         @info "Generating CI plots"
         include("user_io/ci_plots.jl")
-        make_ci_plots([atmos_sim.integrator.p.output_dir], dir_paths.artifacts)
+        make_ci_plots(atmos_sim.integrator.p.output_dir, dir_paths.artifacts)
     end
 
     ## plot all model states and coupler fields (useful for debugging)
@@ -924,4 +918,6 @@ if ClimaComms.iamroot(comms_ctx)
         rm(dir_paths.output; recursive = true, force = true) #hide
     end #hide
 
+    ## close all AMIP diagnostics file writers
+    mode_name == "amip" && map(diag -> close(diag.output_writer), amip_diags_handler.scheduled_diagnostics)
 end
diff --git a/experiments/ClimaEarth/run_cloudless_aquaplanet.jl b/experiments/ClimaEarth/run_cloudless_aquaplanet.jl
@@ -25,15 +25,7 @@ import ClimaCore as CC
 # ## Coupler specific imports
 import ClimaCoupler
 import ClimaCoupler:
-    ConservationChecker,
-    Checkpointer,
-    Diagnostics,
-    FieldExchanger,
-    FluxCalculator,
-    Interfacer,
-    Regridder,
-    TimeManager,
-    Utilities
+    ConservationChecker, Checkpointer, FieldExchanger, FluxCalculator, Interfacer, Regridder, TimeManager, Utilities
 
 pkg_dir = pkgdir(ClimaCoupler)
 
@@ -46,7 +38,6 @@ include("components/atmosphere/climaatmos.jl")
 include("components/ocean/slab_ocean.jl")
 
 ## helpers for user-specified IO
-include("user_io/user_diagnostics.jl")
 include("user_io/user_logging.jl")
 
 include("user_io/io_helpers.jl")
@@ -267,11 +258,11 @@ cs = Interfacer.CoupledSimulation{FT}(
     Δt_cpl,
     model_sims,
     (;), # mode_specifics
-    (), # coupler diagnostics
     callbacks,
     dir_paths,
     turbulent_fluxes,
     thermo_params,
+    nothing, # amip_diags_handler
 );
 
 #=