adding ci plots

config/model_configs/slabplanet_atmos_diags.yml

@@ -1,5 +1,6 @@
 anim: true
 apply_limiter: false
+ci_plots: true
 dt: "200secs"
 dt_cpl: 200
 dt_save_to_sol: "9days"
@@ -9,11 +10,13 @@ job_id: "slabplanet_atmos_diags"
 mode_name: "slabplanet"
 moist: "equil"
 mono_surface: true
+output_default_diagnostics: false
 precip_model: "0M"
 rad: "gray"
 run_name: "slabplanet_atmos_diags"
 t_end: "10days"
 vert_diff: "true"
 diagnostics:
-  - short_name: [mse, lr, ediff]
+  - short_name: [mse, lr, ediff, ts]
+    reduction_time: average
     period: 1days

config/model_configs/slabplanet_default.yml

@@ -14,3 +14,8 @@ rad: "gray"
 run_name: "slabplanet_default"
 t_end: "10days"
 vert_diff: "true"
+output_default_diagnostics: false
+diagnostics:
+ - short_name: [mse, lr, ediff, hfes, evspsbl, ts]
+   period: 1days
+

experiments/AMIP/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.2"
 manifest_format = "2.0"
-project_hash = "b3a2000766f3d77386965c6863b17e8927e1f502"
+project_hash = "c00c8204c76db2774e82408096e51d91be9ef6bf"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "016833eb52ba2d6bea9fcb50ca295980e728ee24"
@@ -388,6 +388,12 @@ git-tree-sha1 = "ded3e0f3e7069f7c807f7b56caff232921bc2f5f"
 uuid = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 version = "0.2.8"
 
+[[deps.ClimaCoreSpectra]]
+deps = ["ClimaCore", "FFTW"]
+git-tree-sha1 = "697b785d474be925987005655a5e5dc21d0cb0d2"
+uuid = "c2caaa1d-32ae-4754-ba0d-80e7561362e9"
+version = "0.1.3"
+
 [[deps.ClimaCoreTempestRemap]]
 deps = ["ClimaComms", "ClimaCore", "CommonDataModel", "Dates", "LinearAlgebra", "NCDatasets", "PkgVersion", "TempestRemap_jll"]
 git-tree-sha1 = "ac11cc8ad2c043ab753d6888c224c7e2f35f42c0"
@@ -1639,6 +1645,12 @@ weakdeps = ["ChainRulesCore"]
     [deps.LinearOperators.extensions]
     LinearOperatorsChainRulesCoreExt = "ChainRulesCore"
 
+[[deps.LittleCMS_jll]]
+deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll"]
+git-tree-sha1 = "08ed30575ffc5651a50d3291beaf94c3e7996e55"
+uuid = "d3a379c0-f9a3-5b72-a4c0-6bf4d2e8af0f"
+version = "2.15.0+0"
+
 [[deps.LogExpFunctions]]
 deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"]
 git-tree-sha1 = "18144f3e9cbe9b15b070288eef858f71b291ce37"
@@ -1950,6 +1962,12 @@ git-tree-sha1 = "a4ca623df1ae99d09bc9868b008262d0c0ac1e4f"
 uuid = "18a262bb-aa17-5467-a713-aee519bc75cb"
 version = "3.1.4+0"
 
+[[deps.OpenJpeg_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Libtiff_jll", "LittleCMS_jll", "libpng_jll"]
+git-tree-sha1 = "8d4c87ffaf09dbdd82bcf8c939843e94dd424df2"
+uuid = "643b3616-a352-519d-856d-80112ee9badc"
+version = "2.5.0+0"
+
 [[deps.OpenLibm_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
@@ -2168,6 +2186,12 @@ git-tree-sha1 = "36d8b4b899628fb92c2749eb488d884a926614d3"
 uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
 version = "1.4.3"
 
+[[deps.Poppler_jll]]
+deps = ["Artifacts", "Cairo_jll", "Fontconfig_jll", "FreeType2_jll", "Glib_jll", "JLLWrappers", "JpegTurbo_jll", "LibCURL_jll", "Libdl", "Libtiff_jll", "OpenJpeg_jll", "libpng_jll"]
+git-tree-sha1 = "a524f03b48f0a90eea898372353e90381ea5ecf4"
+uuid = "9c32591e-4766-534b-9725-b71a8799265b"
+version = "23.12.0+1"
+
 [[deps.PositiveFactorizations]]
 deps = ["LinearAlgebra"]
 git-tree-sha1 = "17275485f373e6673f7e7f97051f703ed5b15b20"

experiments/AMIP/Project.toml

@@ -8,6 +8,7 @@ ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaCorePlots = "cf7c7e5a-b407-4c48-9047-11a94a308626"
+ClimaCoreSpectra = "c2caaa1d-32ae-4754-ba0d-80e7561362e9"
 ClimaCoreTempestRemap = "d934ef94-cdd4-4710-83d6-720549644b70"
 ClimaCoupler = "4ade58fe-a8da-486c-bd89-46df092ec0c7"
 ClimaLand = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
@@ -33,6 +34,7 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Poppler_jll = "9c32591e-4766-534b-9725-b71a8799265b"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

experiments/AMIP/cli_options.jl

@@ -18,6 +18,10 @@ function argparse_settings()
         help = "Boolean flag indicating whether to check energy conservation"
         arg_type = Bool
         default = false
+        "--ci_plots"
+        help = "Boolean flag indicating whether to make CI plots"
+        arg_type = Bool
+        default = false
         "--conservation_softfail"
         help = "Boolean flag indicating whether to soft fail on conservation errors"
         arg_type = Bool

experiments/AMIP/coupler_driver.jl

@@ -882,6 +882,13 @@ if ClimaComms.iamroot(comms_ctx)
         end
     end
 
+    ## ci plots
+    if config_dict["ci_plots"]
+        @info "Generating CI plots"
+        include("user_io/ci_plots.jl")
+        make_plots(Val(:general_ci_plots), [joinpath(COUPLER_OUTPUT_DIR, "clima_atmos")], COUPLER_ARTIFACTS_DIR)
+    end
+
     if isinteractive()
         ## clean up for interactive runs, retain all output otherwise
         rm(COUPLER_OUTPUT_DIR; recursive = true, force = true)

experiments/AMIP/user_io/ci_plots.jl

@@ -0,0 +1,196 @@
+# this file follows the ClimaAtmos.Diagnostics and ci_plots interfaces
+import CairoMakie
+import CairoMakie.Makie
+import ClimaAnalysis as CAN
+using Poppler_jll: pdfunite
+import Base.Filesystem
+
+const LARGE_NUM = typemax(Int)
+const LAST_SNAP = LARGE_NUM
+const FIRST_SNAP = -LARGE_NUM
+const BOTTOM_LVL = -LARGE_NUM
+const TOP_LVL = LARGE_NUM
+
+function Makie.get_tickvalues(yticks::Int, ymin, ymax)
+    return range(max(ymin, 0), ymax, yticks)
+end
+
+YLINEARSCALE = Dict(:axis => CAN.Utils.kwargs(dim_on_y = true, yticks = 10, ytickformat = "{:.3e}"))
+
+long_name(var) = var.attributes["long_name"]
+short_name(var) = var.attributes["short_name"]
+
+"""
+    make_plots_generic(
+        file_path::Union{<:AbstractString, Vector{<:AbstractString}},
+        plot_path,
+        vars,
+        args...;
+        plot_fn = nothing,
+        output_name = "summary",
+        summary_files = String[],
+        MAX_NUM_COLS = 1,
+        MAX_NUM_ROWS = min(4, length(vars)),
+        kwargs...,
+    )
+Create plots for each variable in `vars` and save them to `plot_path`. The number of plots per
+page is determined by `MAX_NUM_COLS` and `MAX_NUM_ROWS`. The `plot_fn` function is used to create the
+plots. If `plot_fn` is not provided, a default plotting function is used. The default plotting function
+is determined by the keyword arguments `kwargs`.
+"""
+function make_plots_generic(
+    file_path::Union{<:AbstractString, Vector{<:AbstractString}},
+    plot_path,
+    vars,
+    args...;
+    plot_fn = nothing,
+    output_name = "summary",
+    summary_files = String[],
+    MAX_NUM_COLS = 1,
+    MAX_NUM_ROWS = min(4, length(vars)),
+    kwargs...,
+)
+    # When file_path is a Vector with multiple elements, this means that this function is
+    # being used to produce a comparison plot. In that case, we modify the output name, and
+    # the number of columns (to match how many simulations we are comparing).
+    is_comparison = file_path isa Vector
+    #
+    # However, we don't want to do this when the vector only contains one element.
+    if is_comparison && length(file_path) == 1
+        # Fallback to the "file_path isa String" case
+        file_path = file_path[1]
+        is_comparison = false
+    end
+
+    if is_comparison
+        MAX_NUM_COLS = length(file_path)
+        plot_path = file_path[1]
+        output_name *= "_comparison"
+    end
+
+    # Default plotting function needs access to kwargs
+    if isnothing(plot_fn)
+        plot_fn = (grid_loc, var) -> CAN.Visualize.plot!(grid_loc, var, args...; kwargs...)
+    end
+
+    MAX_PLOTS_PER_PAGE = MAX_NUM_ROWS * MAX_NUM_COLS
+    vars_left_to_plot = length(vars)
+
+    # Define fig, grid, and grid_pos, used below. (Needed for scope)
+    function makefig()
+        fig = CairoMakie.Figure(; size = (900, 300 * MAX_NUM_ROWS))
+        if is_comparison
+            for (col, path) in (file_path)
+                # CairoMakie seems to use this Label to determine the width of the figure.
+                # Here we normalize the length so that all the columns have the same width.
+                LABEL_LENGTH = 40
+                path = convert(Vector{Float64}, path)
+                normalized_path = lpad(path, LABEL_LENGTH + 1, " ")[(end - LABEL_LENGTH):end]
+
+                CairoMakie.Label(fig[0, col], path)
+            end
+        end
+        return fig
+    end
+
+    # Standardizes grid layout
+    gridlayout() =
+        map(1:MAX_PLOTS_PER_PAGE) do i
+            row = mod(div(i - 1, MAX_NUM_COLS), MAX_NUM_ROWS) + 1
+            col = mod(i - 1, MAX_NUM_COLS) + 1
+            return fig[row, col] = CairoMakie.GridLayout()
+        end
+
+    fig = makefig()
+    grid = gridlayout()
+    page = 1
+    grid_pos = 1
+
+    for var in vars
+        if grid_pos > MAX_PLOTS_PER_PAGE
+            fig = makefig()
+            grid = gridlayout()
+            grid_pos = 1
+        end
+
+        plot_fn(grid[grid_pos], var)
+        grid_pos += 1
+
+        # Flush current page
+        if grid_pos > min(MAX_PLOTS_PER_PAGE, vars_left_to_plot)
+            file_path = joinpath(plot_path, "$(output_name)_$page.pdf")
+            CairoMakie.resize_to_layout!(fig)
+            CairoMakie.save(file_path, fig)
+            push!(summary_files, file_path)
+            vars_left_to_plot -= MAX_PLOTS_PER_PAGE
+            page += 1
+        end
+    end
+
+    # Save plots
+    output_file = joinpath(plot_path, "$(output_name).pdf")
+
+    pdfunite() do unite
+        run(Cmd([unite, summary_files..., output_file]))
+    end
+
+    # Cleanup
+    Filesystem.rm.(summary_files, force = true)
+    return output_file
+end
+
+function map_comparison(func, simdirs, args)
+    return vcat([[func(simdir, arg) for simdir in simdirs] for arg in args]...)
+end
+
+"""
+    make_plots(
+        ::Union{Val{:general_ci_plots}},
+        output_paths::Vector{<:AbstractString},
+        plot_path::AbstractString;
+        reduction::String = "average",
+    )
+Create plots for the general CI diagnostics. The plots are saved to `plot_path`.
+This is the default plotting function for the CI diagnostics and it can be extended
+to include additional diagnostics.
+The `reduction` keyword argument should be consistent with the reduction used to save the diagnostics.
+"""
+function make_plots(
+    ::Union{Val{:general_ci_plots}},
+    output_paths::Vector{<:AbstractString},
+    plot_path::AbstractString;
+    reduction::String = "average",
+)
+    simdirs = CAN.SimDir.(output_paths)
+
+    # Default output diagnostics
+    short_names_3D = ["mse", "lr", "ediff"]
+    short_names_2D = ["ts"]
+
+    available_periods = CAN.available_periods(simdirs[1]; short_name = short_names_3D[1], reduction)
+    period = ""
+    if "10d" in available_periods
+        period = "10d"
+    elseif "1d" in available_periods
+        period = "1d"
+    elseif "12h" in available_periods
+        period = "12h"
+    end
+
+    # Creates diagnostics vector
+    # 3D fields are zonally averaged platted onf the lat-z plane
+    # 2D fields are plotted on the lon-lat plane
+    vars_3D = map_comparison(simdirs, short_names_3D) do simdir, short_name
+        get(simdir; short_name, reduction, period) |> CAN.average_lon
+    end
+
+    available_periods = CAN.available_periods(simdirs[1]; short_name = short_names_2D[1], reduction)
+
+    vars_2D = map_comparison(simdirs, short_names_2D) do simdir, short_name
+        get(simdir; short_name, reduction, period)
+    end
+
+    make_plots_generic(output_paths, plot_path, vars_3D, time = LAST_SNAP, more_kwargs = YLINEARSCALE)
+    make_plots_generic(output_paths, plot_path, vars_2D, time = LAST_SNAP, output_name = "summary_2D")
+
+end