Add SoilCanopyModel diagnostics

This PR finalizes the addition of SoilCanopyModel diagnostics.
It also adds a macro to generate land compute functions,
an optional argument for period averaging for default diagnostics,
an optional argument for ":long" (all) or ":short" (a few essentials)
for default diagnostics.
Diagnostics are now used in global_soil_canopy.jl.
EnergyHydrology model diagnostics_compute methods have been added.

Co-authored-by: AlexisRenchon <a.renchon@gmail.com>
Co-authored-by: SBozzolo <gbozzola@caltech.edu>
Co-authored-by: kmdeck <kdeck@caltech.edu>
Co-authored-by: braghiere <renato.braghiere@gmail.com>

.buildkite/Project.toml

@@ -42,4 +42,5 @@ cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [compat]
 ClimaTimeSteppers = "0.7"
+ClimaAnalysis = "0.5.7"
 Statistics = "1"

.buildkite/pipeline.yml

@@ -119,7 +119,7 @@ steps:
 
       - label: "Global Run CPU"
         command: "julia --color=yes --project=.buildkite experiments/integrated/global/global_soil_canopy.jl"
-        artifact_paths: "experiments/integrated/global/plots/*png"
+        artifact_paths: "experiments/integrated/global/output_active/*png"
         agents:
           slurm_mem: 16G
 
@@ -147,7 +147,7 @@ steps:
         command: "julia --color=yes --project=.buildkite experiments/standalone/Bucket/global_bucket_function.jl"
         artifact_paths:
           - "experiments/standalone/Bucket/artifacts/*cpu*"
-          - "experiments/global_bucket_function/output_active/*.png"
+          - "experiments/standalone/Bucket/artifacts_function/output_active/*.png"
 
       - label: "Global Bucket on CPU (static map albedo)"
         key: "bucket_era5_cpu"

docs/src/diagnostics/make_diagnostic_table.jl

@@ -16,7 +16,14 @@ for d in values(CL.Diagnostics.ALL_DIAGNOSTICS)
     push!(comments, d.comments)
     push!(standard_names, d.standard_name)
 end
-data = hcat(short_names, long_names, units, comments, standard_names)
+i = sortperm(short_names) # indices of short_names sorted alphabetically
+data = hcat(
+    short_names[i],
+    long_names[i],
+    units[i],
+    comments[i],
+    standard_names[i],
+)
 pretty_table(
     data;
     autowrap = true,

docs/src/diagnostics/users_diagnostics.md

@@ -58,6 +58,10 @@ sol = SciMLBase.solve(prob, ode_algo; dt = Δt, callback = diag_cb)
 
 Your diagnostics have now been written in netcdf files in your output folder.
 
+Note that by default, `default_diagnostics` assign two optional kwargs: `output_vars = :long` and `average_period` = :daily.
+`output_vars = :long` will write all available diagnostics, whereas `output_vars = :short` will only write essentials diagnostics. 
+`average_period` defines the period over which diagnostics are averaged, it can be set to `:hourly`, `:daily` and `:monthly`.  
+
 # Custom Diagnostics
 
 When defining a custom diagnostic, follow these steps:

experiments/integrated/global/global_soil_canopy.jl

@@ -24,11 +24,16 @@ using Statistics
 using Dates
 import NCDatasets
 
+using ClimaDiagnostics
+using ClimaAnalysis
+import ClimaAnalysis.Visualize as viz
+using ClimaUtilities
+
 regridder_type = :InterpolationsRegridder
 extrapolation_bc =
     (Interpolations.Periodic(), Interpolations.Flat(), Interpolations.Flat())
 context = ClimaComms.context()
-outdir = joinpath(pkgdir(ClimaLand), "experiments/integrated/global/plots")
+outdir = joinpath(pkgdir(ClimaLand), "experiments/integrated/global")
 !ispath(outdir) && mkpath(outdir)
 
 device_suffix =
@@ -361,76 +366,37 @@ prob = SciMLBase.ODEProblem(
     p,
 )
 
-saveat = [t0, tf]
-sv = (;
-    t = Array{Float64}(undef, length(saveat)),
-    saveval = Array{NamedTuple}(undef, length(saveat)),
-)
-saving_cb = ClimaLand.NonInterpSavingCallback(sv, saveat)
+# ClimaDiagnostics
+output_dir = ClimaUtilities.OutputPathGenerator.generate_output_path(outdir)
+
+nc_writer = ClimaDiagnostics.Writers.NetCDFWriter(subsurface_space, output_dir)
+
+diags = ClimaLand.CLD.default_diagnostics(land, t0; output_writer = nc_writer)
+
+diagnostic_handler =
+    ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0; dt = dt)
+
+diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler)
+
 updateat = Array(t0:(3600 * 3):tf)
 drivers = ClimaLand.get_drivers(land)
 updatefunc = ClimaLand.make_update_drivers(drivers)
 driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
-@time sol = SciMLBase.solve(
+
+sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
     prob,
     ode_algo;
     dt = dt,
-    callback = cb,
-    adaptive = false,
-    saveat = saveat,
+    callback = SciMLBase.CallbackSet(driver_cb, diag_cb),
 )
 
-if device_suffix == "cpu"
-    longpts = range(-180.0, 180.0, 101)
-    latpts = range(-90.0, 90.0, 101)
-    hcoords = [
-        ClimaCore.Geometry.LatLongPoint(lat, long) for long in longpts,
-        lat in reverse(latpts)
-    ]
-    remapper = ClimaCore.Remapping.Remapper(surface_space, hcoords)
-    S = ClimaLand.Domains.top_center_to_surface(
-        (sol.u[end].soil.ϑ_l .- θ_r) ./ (ν .- θ_r),
-    )
-    S_ice = ClimaLand.Domains.top_center_to_surface(sol.u[end].soil.θ_i ./ ν)
-    T_soil = ClimaLand.Domains.top_center_to_surface(sv.saveval[end].soil.T)
-    SW = sv.saveval[end].drivers.SW_d
-
-    GPP = sv.saveval[end].canopy.photosynthesis.GPP .* 1e6
-    T_canopy = sol.u[end].canopy.energy.T
-    fields = [S, S_ice, T_soil, GPP, T_canopy, SW]
-    titles = [
-        "Effective saturation",
-        "Effective ice saturation",
-        "Temperature (K) - Soil",
-        "GPP",
-        "Temperature (K) - Canopy",
-        "Incident SW",
-    ]
-    plotnames = ["S", "Sice", "temp", "gpp", "temp_canopy", "sw"]
-    mask_remap = ClimaCore.Remapping.interpolate(remapper, soil_params_mask_sfc)
-    for (id, x) in enumerate(fields)
-        title = titles[id]
-        plotname = plotnames[id]
-        x_remap = ClimaCore.Remapping.interpolate(remapper, x)
-
-        fig = Figure(size = (600, 400))
-        ax = Axis(
-            fig[1, 1],
-            xlabel = "Longitude",
-            ylabel = "Latitude",
-            title = title,
-        )
-        clims = extrema(x_remap)
-        CairoMakie.heatmap!(
-            ax,
-            longpts,
-            latpts,
-            oceans_to_value.(x_remap, mask_remap, 0.0),
-            colorrange = clims,
-        )
-        Colorbar(fig[:, end + 1], colorrange = clims)
-        outfile = joinpath(outdir, "$plotname.png")
-        CairoMakie.save(outfile, fig)
-    end
+# ClimaAnalysis
+simdir = ClimaAnalysis.SimDir(outdir)
+
+for short_name in ClimaAnalysis.available_vars(simdir)
+    var = get(simdir; short_name)
+    fig = CairoMakie.Figure(size = (800, 600))
+    kwargs = ClimaAnalysis.has_altitude(var) ? Dict(:z => 1) : Dict()
+    viz.plot!(fig, var, lat = 0; kwargs...)
+    CairoMakie.save(joinpath(output_dir, "$short_name.png"), fig)
 end

experiments/long_runs/land.jl

@@ -602,7 +602,7 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 15))
         land,
         t0;
         output_writer = nc_writer,
-        output_vars = :short,
+        output_vars = :long,
     )
 
     diagnostic_handler =
@@ -646,14 +646,14 @@ setup_and_solve_problem(; greet = true);
 # read in diagnostics and make some plots!
 #### ClimaAnalysis ####
 simdir = ClimaAnalysis.SimDir(outdir)
-short_names_2D = ["gpp", "ct", "slw", "si"]
-times = 0.0:(60.0 * 60.0 * 12):(60.0 * 60.0 * 24 * 7)
-for t in times
-    for short_name in short_names_2D
-        var = get(simdir; short_name)
+short_names = ["gpp", "swc", "si", "sie"]
+for short_name in short_names
+    var = get(simdir; short_name)
+    times = var.dims["time"]
+    for t in times
         fig = CairoMakie.Figure(size = (800, 600))
-        kwargs = short_name in short_names_2D ? Dict() : Dict(:z => 1)
+        kwargs = ClimaAnalysis.has_altitude(var) ? Dict(:z => 1) : Dict()
         viz.plot!(fig, var, time = t; kwargs...)
-        CairoMakie.save(joinpath(root_path, "$short_name $t.png"), fig)
+        CairoMakie.save(joinpath(root_path, "$(short_name)_$t.png"), fig)
     end
 end

experiments/long_runs/soil.jl

@@ -464,14 +464,15 @@ setup_and_solve_problem(; greet = true);
 # read in diagnostics and make some plots!
 #### ClimaAnalysis ####
 simdir = ClimaAnalysis.SimDir(outdir)
-short_names_2D = ["slw", "si", "tsoil"]
-times = 0.0:(60.0 * 60.0 * 24 * 20):(60.0 * 60.0 * 24 * 60)
-for t in times
-    for short_name in short_names_2D
+short_names = ["swc", "si", "sie"]
+for short_name in short_names
+    var = get(simdir; short_name)
+    times = var.dims["time"]
+    for t in times
         var = get(simdir; short_name)
         fig = CairoMakie.Figure(size = (800, 600))
-        kwargs = short_name in short_names_2D ? Dict() : Dict(:z => 1)
+        kwargs = ClimaAnalysis.has_altitude(var) ? Dict(:z => 1) : Dict()
         viz.plot!(fig, var, time = t; kwargs...)
-        CairoMakie.save(joinpath(root_path, "$short_name $t.png"), fig)
+        CairoMakie.save(joinpath(root_path, "$(short_name)_$t.png"), fig)
     end
 end

experiments/standalone/Bucket/global_bucket_function.jl

@@ -162,9 +162,7 @@ prob = SciMLBase.ODEProblem(
 );
 
 # ClimaDiagnostics
-base_output_dir = "global_bucket_function/"
-output_dir =
-    ClimaUtilities.OutputPathGenerator.generate_output_path(base_output_dir)
+output_dir = ClimaUtilities.OutputPathGenerator.generate_output_path(outdir)
 
 space = bucket_domain.space.subsurface