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>
CliMA · Aug 17, 2024 · 0fd14bf · 0fd14bf
1 parent 7c9bc04
commit 0fd14bf
Show file tree

Hide file tree

Showing 16 changed files with 915 additions and 1,214 deletions.
diff --git a/.buildkite/Manifest.toml b/.buildkite/Manifest.toml
diff --git a/.buildkite/Project.toml b/.buildkite/Project.toml
@@ -7,8 +7,8 @@ CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ClimaAnalysis = "29b5916a-a76c-4e73-9657-3c8fd22e65e6"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
-ClimaDiagnostics = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
 ClimaCoreMakie = "908f55d8-4145-4867-9c14-5dad1a479e4d"
+ClimaDiagnostics = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
 ClimaLand = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
@@ -18,6 +18,7 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Formatting = "59287772-0a20-5a39-b81b-1366585eb4c0"
+GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6"
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
@@ -41,5 +42,6 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [compat]
+ClimaAnalysis = "0.5.7"
 ClimaTimeSteppers = "0.7"
 Statistics = "1"
diff --git a/.buildkite/pipeline.yml b/.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"

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
@@ -24,7 +24,7 @@ output_dir = ClimaUtilities.OutputPathGenerator.generate_output_path("base_outpu
 
 2. define a space
 
-Your diagnostics will be written in time and space. These may be defined in your model, but usually land model space is a sphere with no vertical dimension. 
+Your diagnostics will be written in time and space. These may be defined in your model, but usually land model space is a sphere with no vertical dimension.
 You may have variables varying with soil depth, and so you will need:
 
 ```
@@ -46,7 +46,7 @@ providing the space and output_dir defined in steps 1. and 2.
 Now that you defined your model and your writter, you can create a callback function to be called when solving your model. For example:
 
 ```
-diags = ClimaLand.CLD.default_diagnostics(model, 1.0; output_writer = nc_writer)
+diags = ClimaLand.default_diagnostics(model, 1.0; output_writer = nc_writer)
 
 diagnostic_handler =
     ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0; dt = Δt)
@@ -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:
@@ -82,7 +86,7 @@ end
         compute! = (out, Y, p, t) -> compute_bowen_ratio!(out, Y, p, t, land_model),
     )
  ```
- - Define how to schedule your variables. For example, you want the seasonal maximum of your variables, where season is defined as 90 days. 
+ - Define how to schedule your variables. For example, you want the seasonal maximum of your variables, where season is defined as 90 days.
  ```
  seasonal_maxs(short_names...; output_writer, t_start) = common_diagnostics(
     90 * 24 * 60 * 60 * one(t_start),

diff --git a/experiments/integrated/global/global_soil_canopy.jl b/experiments/integrated/global/global_soil_canopy.jl
@@ -19,16 +19,22 @@ using ClimaLand.Canopy
 import ClimaLand
 import ClimaLand.Parameters as LP
 
-using CairoMakie
+import CairoMakie
+import GeoMakie
 using Statistics
 using Dates
 import NCDatasets
 
+import ClimaDiagnostics
+import ClimaAnalysis
+import ClimaAnalysis.Visualize as viz
+import 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 +367,46 @@ prob = SciMLBase.ODEProblem(
     p,
 )
 
-saveat = [t0, tf]
-sv = (;
-    t = Array{Float64}(undef, length(saveat)),
-    saveval = Array{NamedTuple}(undef, length(saveat)),
+# ClimaDiagnostics
+output_dir = ClimaUtilities.OutputPathGenerator.generate_output_path(outdir)
+
+nc_writer = ClimaDiagnostics.Writers.NetCDFWriter(subsurface_space, output_dir)
+
+diags = ClimaLand.default_diagnostics(
+    land,
+    t0;
+    output_writer = nc_writer,
+    average_period = :hourly,
 )
-saving_cb = ClimaLand.NonInterpSavingCallback(sv, saveat)
+
+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(output_dir)
+
+for short_name in ClimaAnalysis.available_vars(simdir)
+    var = get(simdir; short_name)
+    times = var.dims["time"]
+    for t in times 
+        fig = CairoMakie.Figure(size = (800, 600))
+        kwargs = ClimaAnalysis.has_altitude(var) ? Dict(:z => 1) : Dict()
+        viz.heatmap2D_on_globe!(fig, ClimaAnalysis.slice(var, time = t; kwargs...))
+        CairoMakie.save(joinpath(output_dir, "$short_name.png"), fig)
+    end 
 end
+
diff --git a/experiments/long_runs/land.jl b/experiments/long_runs/land.jl
@@ -23,10 +23,10 @@ using ClimaUtilities.ClimaArtifacts
 import Interpolations
 using Insolation
 
-using ClimaDiagnostics
-using ClimaAnalysis
+import ClimaDiagnostics
+import ClimaAnalysis
 import ClimaAnalysis.Visualize as viz
-using ClimaUtilities
+import ClimaUtilities
 
 import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput
 import ClimaUtilities.SpaceVaryingInputs: SpaceVaryingInput
@@ -42,6 +42,7 @@ import ClimaLand.Parameters as LP
 
 using Statistics
 using CairoMakie
+import GeoMakie
 using Dates
 import NCDatasets
 
@@ -598,11 +599,11 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 15))
 
     nc_writer = ClimaDiagnostics.Writers.NetCDFWriter(subsurface_space, outdir)
 
-    diags = ClimaLand.CLD.default_diagnostics(
+    diags = ClimaLand.default_diagnostics(
         land,
         t0;
         output_writer = nc_writer,
-        output_vars = :short,
+        output_vars = :long,
     )
 
     diagnostic_handler =
@@ -646,14 +647,17 @@ 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 = ClimaAnalysis.times(var)
+    for t in times
         fig = CairoMakie.Figure(size = (800, 600))
-        kwargs = short_name in short_names_2D ? Dict() : Dict(:z => 1)
-        viz.plot!(fig, var, time = t; kwargs...)
-        CairoMakie.save(joinpath(root_path, "$short_name $t.png"), fig)
+        kwargs = ClimaAnalysis.has_altitude(var) ? Dict(:z => 1) : Dict()
+        viz.heatmap2D_on_globe!(
+            fig,
+            ClimaAnalysis.slice(var, time = t; kwargs...),
+        )
+        CairoMakie.save(joinpath(root_path, "$(short_name)_$t.png"), fig)
     end
 end
diff --git a/experiments/long_runs/soil.jl b/experiments/long_runs/soil.jl
@@ -420,8 +420,7 @@ function setup_prob(t0, tf, Δt; outdir = outdir, nelements = (101, 15))
 
     nc_writer = ClimaDiagnostics.Writers.NetCDFWriter(subsurface_space, outdir)
 
-    diags =
-        ClimaLand.CLD.default_diagnostics(soil, t0; output_writer = nc_writer)
+    diags = ClimaLand.default_diagnostics(soil, t0; output_writer = nc_writer)
 
     diagnostic_handler =
         ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0; dt = Δt)
@@ -464,14 +463,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 = ClimaAnalysis.times(var)
+    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
diff --git a/experiments/standalone/Bucket/global_bucket_function.jl b/experiments/standalone/Bucket/global_bucket_function.jl
@@ -162,15 +162,13 @@ 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
 
 nc_writer = ClimaDiagnostics.Writers.NetCDFWriter(space, output_dir)
 
-diags = ClimaLand.CLD.default_diagnostics(model, t0; output_writer = nc_writer)
+diags = ClimaLand.default_diagnostics(model, t0; output_writer = nc_writer)
 
 diagnostic_handler =
     ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0; dt = Δt)
@@ -212,7 +210,7 @@ for short_name in vcat(short_names_2D..., short_names_3D...)
     var = get(simdir; short_name)
     fig = CairoMakie.Figure(size = (800, 600))
     kwargs = short_name in short_names_2D ? Dict() : Dict(:z => 1)
-    viz.plot!(fig, var, lat = 0; kwargs...)
+    viz.plot!(fig, var, lon = 0, lat = 0; kwargs...)
     CairoMakie.save(joinpath(output_dir, "$short_name.png"), fig)
 end
 

diff --git a/src/ClimaLand.jl b/src/ClimaLand.jl
@@ -320,6 +320,6 @@ include("integrated/soil_canopy_model.jl")
 
 # Diagnostics
 include(joinpath("diagnostics", "Diagnostics.jl"))
-import .Diagnostics as CLD # ClimaLand Diagnostics
+import .Diagnostics: default_diagnostics
 
 end