Group EDMF vars before slicing, other fixes

post_processing/ci_plots.jl

@@ -4,7 +4,6 @@ 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
 
@@ -63,6 +62,42 @@ YLOGSCALE = Dict(
 long_name(var) = var.attributes["long_name"]
 short_name(var) = var.attributes["short_name"]
 
+"""
+    parse_var_attributes(var)
+
+Takes in an OutputVar and parses some of its attributes into a short, informative string.
+Used to generate unique titles when the same var is being plotted for several times/locations.
+This could be extended to parse more attributes.
+
+For example, the sample attributes:
+attributes = Dict(
+    "units" => "%",
+    "short_name" => "cl",
+    "slice_y" => "0.0",
+    "long_name" => "Cloud fraction, Instantaneous x = 0.0 m y = 0.0 m",
+    "slice_y_units" => "m",
+    "slice_x_units" => "m",
+    "comments" => "",
+    "slice_x" => "0.0",
+)
+will be parsed into "cl, x = 0.0, y = 0.0"
+"""
+function parse_var_attributes(var)
+    MISSING_STR = "MISSING_ATTRIBUTE"
+    attr = var.attributes
+    name = replace(short_name(var), "up" => "")
+
+    attributes = ["slice_lat", "slice_lon", "slice_x", "slice_y", "slice_time"]
+    info = [
+        replace(key, "slice_" => "") * " = " * get(attr, key, MISSING_STR)
+        for key in attributes
+    ]
+    # Filter out missing entries
+    info = filter(x -> !occursin(MISSING_STR, x), [name, info...])
+
+    return join(info, ", ")
+end
+
 function make_plots_generic(
     output_path,
     vars,
@@ -623,7 +658,7 @@ EDMFBoxPlots = Union{
 }
 
 """
-    plot_edmf_vert_profile!(fig, p_loc, var_group)
+    plot_edmf_vert_profile!(grid_loc, var_group)
 
 Helper function for `make_plots_generic`. Takes a list of variables and plots
 them on the same axis.
@@ -632,48 +667,66 @@ function plot_edmf_vert_profile!(grid_loc, var_group)
     z = var_group[1].dims["z"]
     units = var_group[1].attributes["units"]
     ax = CairoMakie.Axis(
-        grid_loc[1,1],
+        grid_loc[1, 1],
         ylabel = "z [$(var_group[1].dim_attributes["z"]["units"])]",
         xlabel = "$(short_name(var_group[1])) [$units]",
-        title = short_var_info(var_group[1])
+        title = parse_var_attributes(var_group[1]),
     )
-
     for var in var_group
         CairoMakie.lines!(ax, var.data, z, label = short_name(var))
     end
     length(var_group) > 1 && Makie.axislegend(ax)
 end
 
+
+"""
+    plot_parsed_attribute_title!(grid_loc, var)
+
+Helper function for `make_plots_generic`. Plots an OutputVar `var`,
+setting the axis title to `parse_var_attributes(var)`
+"""
+plot_parsed_attribute_title!(grid_loc, var) = viz.plot!(
+    grid_loc,
+    var;
+    more_kwargs = Dict(:axis => ca_kwargs(title = parse_var_attributes(var))),
+)
+
 """
-    group_edmf_vars(vars)
+    pair_edmf_names(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.
+Groups updraft and gridmean EDMF short names into tuples. 
+Matches on the same variable short name with the suffix "up".
+This assumes that the updraft variable name is the same as the corresponding
+gridmean variable with the suffix "up".
 """
-function group_edmf_vars(vars)
-    grouped_vars = OrderedDict{String, Vector{ClimaAnalysis.OutputVar}}()
-    for var in sort(vars, by = short_var_info)
-        suffix = short_var_info(var)
-        if haskey(grouped_vars, suffix)
-            push!(grouped_vars[suffix], var)
+function pair_edmf_names(short_names)
+    grouped_vars = Any[]
+    short_names_to_be_processed = Set(short_names)
+
+    for name in short_names
+        # If we have already visited this name, go to the next one
+        name in short_names_to_be_processed || continue
+
+        # First, check if we have the pair of variables
+        # We normalize the name to the gridmean version (base_name)
+        # So, if we are visiting "va" or "vaup", we end up with 
+        # base_name = "va" and up_name = "vaup"
+        base_name = replace(name, "up" => "")
+        up_name = base_name * "up"
+
+        if base_name in short_names_to_be_processed &&
+           up_name in short_names_to_be_processed
+            # Gridmean and updraft are available
+            tuple_to_be_added = (base_name, up_name)
         else
-            grouped_vars[suffix] = [var]
+            # Only single var (updraft OR gridmean) is available
+            tuple_to_be_added = (name,)
         end
-    end
-    return grouped_vars
-end
 
-function short_var_info(var)
-    attr = var.attributes
-    coords = if haskey(attr, "slice_lat")
-        "lat = $(attr["slice_lat"]), lon = $(attr["slice_lon"])"
-    else
-        "x = $(attr["slice_x"]), y = $(attr["slice_y"])"
+        foreach(n -> delete!(short_names_to_be_processed, n), tuple_to_be_added)
+        push!(grouped_vars, tuple_to_be_added)
     end
-    name = replace(short_name(var), "up" => "")
-    time = "t = " * attr["slice_time"]
-    return join([name, coords, time], ", ")
+    return grouped_vars
 end
 
 function make_plots(::EDMFBoxPlots, simulation_path)
@@ -704,30 +757,34 @@ function make_plots(::EDMFBoxPlots, simulation_path)
     reduction = "inst"
     period = "30m"
 
-    vars = [
-        get(simdir; short_name, reduction, period) for short_name in short_names
+    short_name_tuples = pair_edmf_names(short_names)
+    var_groups_zt = [
+        (
+            slice(get(simdir; short_name, reduction, period), x = 0.0, y = 0.0) for short_name in var_names
+        ) for var_names in short_name_tuples
+    ]
+    var_groups_z = [
+        ([slice(v, time = LAST_SNAP) for v in group]...,) for
+        group in var_groups_zt
     ]
-    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]
-    grouped_vars = group_edmf_vars(vars_z)
 
     tmp_file = make_plots_generic(
         simulation_path,
         output_name = "tmp",
-        values(grouped_vars);
+        var_groups_z;
         plot_fn = plot_edmf_vert_profile!,
         MAX_NUM_COLS = 2,
         MAX_NUM_ROWS = 4,
     )
 
     make_plots_generic(
         simulation_path,
-        vars_zt,
+        vcat((var_groups_zt...)...),
+        plot_fn = plot_parsed_attribute_title!,
         summary_files = [tmp_file],
         MAX_NUM_COLS = 2,
         MAX_NUM_ROWS = 4,
     )
-
 end
 
 EDMFSpherePlots =
@@ -736,32 +793,41 @@ EDMFSpherePlots =
 function make_plots(::EDMFSpherePlots, simulation_path)
     simdir = SimDir(simulation_path)
 
-    short_names = ["ua", "wa", "thetaa", "taup", "haup", "waup", "tke", "arup"]
+    short_names =
+        ["ua", "wa", "waup", "thetaa", "ta", "taup", "haup", "tke", "arup"]
     reduction = "average"
-    vars = [get(simdir; short_name, reduction) for short_name in short_names]
-    vars_zt0_0 = [slice(var, lon = 0.0, lat = 0.0) for var in vars]
-    vars_zt30_0 = [slice(var, lon = 0.0, lat = 30.0) for var in vars]
-    vars_zt60_0 = [slice(var, lon = 0.0, lat = 60.0) for var in vars]
-    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)
+    period = "1h"
+    latitudes = [0.0, 30.0, 60.0, 90.0]
+
+    short_name_tuples = pair_edmf_names(short_names)
+    var_groups_zt = [
+        (
+            slice(
+                get(simdir; short_name, reduction, period),
+                lon = 0.0,
+                lat = lat,
+            ) for short_name in var_names
+        ) for lat in latitudes, var_names in short_name_tuples
+    ]
+    var_groups_z = [
+        ([slice(v, time = LAST_SNAP) for v in group]...,) for
+        group in var_groups_zt
+    ]
 
     tmp_file = make_plots_generic(
         simulation_path,
         output_name = "tmp",
-        values(grouped_vars);
+        var_groups_z;
         plot_fn = plot_edmf_vert_profile!,
         MAX_NUM_COLS = 2,
         MAX_NUM_ROWS = 4,
     )
-
     make_plots_generic(
         simulation_path,
-        vars_zt,
+        vcat((var_groups_zt...)...),
+        plot_fn = plot_parsed_attribute_title!,
         summary_files = [tmp_file],
         MAX_NUM_COLS = 2,
         MAX_NUM_ROWS = 4,
-        more_kwargs = Dict(:axis => ca_kwargs(title = var -> var.attributes["short_name"]))
     )
 end