Refactor make_plots_generic, customize EDMF plot

post_processing/ci_plots.jl

@@ -4,6 +4,7 @@ import ClimaAnalysis
 import ClimaAnalysis: Visualize as viz
 import ClimaAnalysis: SimDir, slice_time, slice
 import ClimaAnalysis.Utils: kwargs as ca_kwargs
+import OrderedCollections: OrderedDict
 
 import ClimaCoreSpectra: power_spectrum_2d
 
@@ -26,7 +27,7 @@ function common_dirname(files::Vector{T}) where {T <: AbstractString}
     return joinpath(split_files[1][1:last_common_dir]...)
 end
 
-function make_plots(sim, simulation_path)
+function make_plots(sim, _)
     @warn "No plot found for $sim"
 end
 
@@ -59,41 +60,57 @@ YLOGSCALE = Dict(
     ),
 )
 
+long_name(var) = var.attributes["long_name"]
+short_name(var) = var.attributes["short_name"]
+
 function make_plots_generic(
     output_path,
     vars,
     args...;
+    plot_fn = nothing,
     output_name = "summary",
+    summary_files = String[],
+    MAX_NUM_COLS = 1,
+    MAX_NUM_ROWS = 4,
     kwargs...,
 )
-    MAX_PLOTS_PER_PAGE = 4
+    # Default plotting function needs access to kwargs
+    if isnothing(plot_fn)
+        plot_fn =
+            (fig, fig_pos, var) -> viz.plot!(fig, var, args...; fig_pos, kwargs...)
+    end
+
+    MAX_PLOTS_PER_PAGE = MAX_NUM_ROWS * MAX_NUM_COLS
     vars_left_to_plot = length(vars)
 
-    # Define fig, and p_loc, used below. (Needed for scope)
-    fig = CairoMakie.Figure(
-        resolution = (900, 300 * min(vars_left_to_plot, MAX_PLOTS_PER_PAGE)),
-    )
-    p_loc = [1, 1]
+    # Define fig, grid, and grid_pos, used below. (Needed for scope)
+    makefig() = CairoMakie.Figure(resolution = (900, 300 * MAX_NUM_ROWS))
+    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
+            fig[row, col] = CairoMakie.GridLayout()
+            return [row, col]
+        end
+
+    fig = makefig()
+    grid = gridlayout()
     page = 1
-    summary_files = String[]
+    grid_pos = 1
 
-    for (var_num, var) in enumerate(vars)
+    for var in vars
         # Create a new page if this is the first plot
-        if mod(var_num, MAX_PLOTS_PER_PAGE) == 1
-            fig = CairoMakie.Figure(
-                resolution = (
-                    900,
-                    300 * min(vars_left_to_plot, MAX_PLOTS_PER_PAGE),
-                ),
-            )
-            p_loc = [1, 1]
+        if grid_pos > MAX_PLOTS_PER_PAGE
+            fig = makefig()
+            grid = gridlayout()
+            grid_pos = 1
         end
-        viz.plot!(fig, var, args...; p_loc, kwargs...)
 
-        p_loc[1] += 1
+        plot_fn(fig, grid[grid_pos], var)
+        grid_pos += 1
 
         # Flush current page
-        if p_loc[1] > min(MAX_PLOTS_PER_PAGE, vars_left_to_plot)
+        if grid_pos > min(MAX_PLOTS_PER_PAGE, vars_left_to_plot)
             file_path = joinpath(output_path, "$(output_name)_$page.pdf")
             CairoMakie.save(file_path, fig)
             push!(summary_files, file_path)
@@ -110,6 +127,7 @@ function make_plots_generic(
 
     # Cleanup
     Filesystem.rm.(summary_files, force = true)
+    return output_file
 end
 
 """
@@ -539,22 +557,95 @@ EDMFBoxPlots = Union{
     Val{:prognostic_edmfx_trmm_column},
 }
 
+"""
+    plot_edmf_vert_profile(fig, p_loc, var_group)
+
+Helper function for `make_plots_generic`. Takes a list of variables and plots
+them on the same axis.
+"""
+function plot_edmf_vert_profile(fig, p_loc, var_group)
+    z = var_group[1].dims["z"]
+    units = var_group[1].attributes["units"]
+
+    ax = CairoMakie.Axis(
+        fig[p_loc...],
+        ylabel = "z [$(var_group[1].dim_attributes["z"]["units"])]",
+        xlabel = "$(short_name(var_group[1])) [$units]",
+    )
+
+    for var in var_group
+        CairoMakie.lines!(ax, var.data, z, label = short_name(var))
+    end
+    length(var_group) > 1 && Makie.axislegend(ax)
+end
+
+"""
+    group_edmf_vars(vars)
+
+Groups pairs of updraft and gridmean EDMF OutputVars into a Dict.
+This is fairly brittle - it just checks that the long name starts with "Updraft"
+and matches on the rest of the long name.
+"""
+function group_edmf_vars(vars)
+    get_longname_suffix(v) =
+        startswith(long_name(v), "Updraft") ? long_name(v)[9:end] : long_name(v)
+    grouped_vars = OrderedDict{String, Vector{ClimaAnalysis.OutputVar}}()
+    for var in sort(vars, by = get_longname_suffix)
+        suffix = get_longname_suffix(var)
+        if haskey(grouped_vars, suffix)
+            push!(grouped_vars[suffix], var)
+        else
+            grouped_vars[suffix] = [var]
+        end
+    end
+    return grouped_vars
+end
+
 function make_plots(::EDMFBoxPlots, simulation_path)
     simdir = SimDir(simulation_path)
 
-    short_names = ["ua", "wa", "thetaa", "taup", "haup", "waup", "tke", "arup"]
+    short_names = [
+        "ua",
+        "wa",
+        "thetaa",
+        "thetaaup",
+        "ta",
+        "taup",
+        "ha",
+        "haup",
+        "waup",
+        "tke",
+        "arup",
+        "hus",
+        "husup",
+        "hur",
+        "hurup",
+        "cl",
+        "clw",
+        "clwup",
+        "cli",
+        "cliup",
+    ]
     reduction = "average"
     period = "10m"
     vars = [
         get(simdir; short_name, reduction, period) for short_name in short_names
     ]
     vars_zt = [slice(var, x = 0.0, y = 0.0) for var in vars]
     vars_z = [slice(var, x = 0.0, y = 0.0, time = LAST_SNAP) for var in vars]
-    make_plots_generic(
+    grouped_vars = group_edmf_vars(vars_z)
+
+    tmp_file = make_plots_generic(
         simulation_path,
-        [vars_zt..., vars_z...],
-        more_kwargs = YLOGSCALE,
+        output_name = "tmp",
+        values(grouped_vars);
+        plot_fn = plot_edmf_vert_profile,
+        MAX_NUM_COLS = 2,
+        MAX_NUM_ROWS = 4,
     )
+
+    make_plots_generic(simulation_path, vars_zt, summary_files = [tmp_file])
+
 end
 
 EDMFSpherePlots =
@@ -575,10 +666,16 @@ function make_plots(::EDMFSpherePlots, simulation_path)
     vars_zt90_0 = [slice(var, lon = 0.0, lat = 90.0) for var in vars]
     vars_zt = [vars_zt0_0..., vars_zt30_0..., vars_zt60_0..., vars_zt90_0...]
     vars_z = [slice(var, time = LAST_SNAP) for var in vars_zt]
+    grouped_vars = group_edmf_vars(vars_z)
 
-    make_plots_generic(
+    tmp_file = make_plots_generic(
         simulation_path,
-        [vars_zt..., vars_z...],
-        more_kwargs = YLOGSCALE,
+        output_name = "tmp",
+        values(grouped_vars);
+        plot_fn = plot_edmf_vert_profile,
+        MAX_NUM_COLS = 2,
+        MAX_NUM_ROWS = 4,
     )
+
+    make_plots_generic(simulation_path, vars_zt, summary_files = [tmp_file])
 end