Add SoilCanopyModel diagnostics

This PR finalizes the addition of SoilCanopyModel diagnostics. It also adds a macro to generate land compute functions, And an optional argument for period averaging for default diagnostics. 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>
CliMA · Aug 16, 2024 · 0d78fc4 · 0d78fc4
1 parent 7c9bc04
commit 0d78fc4
Show file tree

Hide file tree

Showing 12 changed files with 767 additions and 1,055 deletions.
diff --git a/.buildkite/Project.toml b/.buildkite/Project.toml
@@ -42,4 +42,5 @@ cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [compat]
 ClimaTimeSteppers = "0.7"
+ClimaAnalysis = "0.5.7"
 Statistics = "1"
diff --git a/docs/src/diagnostics/make_diagnostic_table.jl b/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,

diff --git a/docs/src/diagnostics/users_diagnostics.md b/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` = :hourly.
+`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:

diff --git a/experiments/integrated/global/global_soil_canopy.jl b/experiments/integrated/global/global_soil_canopy.jl
@@ -24,6 +24,11 @@ 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())
@@ -361,76 +366,90 @@ 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
+base_output_dir = "global_soil_canopy_function/"
+output_dir =
+    ClimaUtilities.OutputPathGenerator.generate_output_path(base_output_dir)
+
+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)
+short_names = [
+    "sif",
+    "ra",
+    "gs",
+    "trans",
+    "crae",
+    "clhf",
+    "cshf",
+    "lai",
+    "msf",
+    "rai",
+    "sai",
+    "gpp",
+    "an",
+    "rd",
+    "vcmax25",
+    "nir",
+    "anir",
+    "rnir",
+    "tnir",
+    "par",
+    "apar",
+    "rpar",
+    "tpar",
+    "lwn",
+    "swn",
+    "soc",
+    "airp",
+    "rain",
+    "lwd",
+    "swd",
+    "snow",
+    "sza",
+    "qsfc",
+    "ws",
+    "infil",
+    "shc",
+    "stc",
+    "swp",
+    "soilrn",
+    "soilrae",
+    "soillhf",
+    "soilshf",
+    "hr",
+    "scd",
+    "scms",
+    "ct",
+    "sco2",
+    "swc",
+    "si",
+    "sie",
+]
+for short_name in short_names
+    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
diff --git a/experiments/long_runs/land.jl b/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,13 +646,64 @@ 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"]
+short_names = [
+    "sif",
+    "ra",
+    "gs",
+    "trans",
+    "crae",
+    "clhf",
+    "cshf",
+    "lai",
+    "msf",
+    "rai",
+    "sai",
+    "gpp",
+    "an",
+    "rd",
+    "vcmax25",
+    "nir",
+    "anir",
+    "rnir",
+    "tnir",
+    "par",
+    "apar",
+    "rpar",
+    "tpar",
+    "lwn",
+    "swn",
+    "soc",
+    "airp",
+    "rain",
+    "lwd",
+    "swd",
+    "snow",
+    "sza",
+    "qsfc",
+    "ws",
+    "infil",
+    "shc",
+    "stc",
+    "swp",
+    "soilrn",
+    "soilrae",
+    "soillhf",
+    "soilshf",
+    "hr",
+    "scd",
+    "scms",
+    "ct",
+    "sco2",
+    "swc",
+    "si",
+    "sie",
+]
 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
+    for short_name in short_names
         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)
     end

diff --git a/experiments/long_runs/soil.jl b/experiments/long_runs/soil.jl
@@ -464,13 +464,13 @@ setup_and_solve_problem(; greet = true);
 # read in diagnostics and make some plots!
 #### ClimaAnalysis ####
 simdir = ClimaAnalysis.SimDir(outdir)
-short_names_2D = ["slw", "si", "tsoil"]
+short_names = ["swc", "si", "sie"]
 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
+    for short_name in short_names
         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)
     end