Group short names before getting variables

post_processing/ci_plots.jl

@@ -673,7 +673,6 @@ function plot_edmf_vert_profile!(grid_loc, var_group)
         xlabel = "$(short_name(var_group[1])) [$units]",
         title = parse_var_attributes(var_group[1]),
     )
-
     for var in var_group
         CairoMakie.lines!(ax, var.data, z, label = short_name(var))
     end
@@ -694,19 +693,24 @@ plot_parsed_attribute_title!(grid_loc, var) = viz.plot!(
 )
 
 """
-    group_edmf_vars(vars)
+    pair_edmf_names(vars)
 
-Groups updraft and gridmean EDMF OutputVars into a Dict.
-Groupings are matched via variable short name, location, and time (if applicable).
+Groups updraft and gridmean EDMF short names into tuples.
 """
-function group_edmf_vars(vars)
-    grouped_vars = OrderedDict{String, Vector{ClimaAnalysis.OutputVar}}()
-    for var in sort(vars, by = parse_var_attributes)
-        info = parse_var_attributes(var)
-        if haskey(grouped_vars, info)
-            push!(grouped_vars[info], var)
+function pair_edmf_names(short_names)
+    grouped_vars = Any[]
+
+    for name in short_names
+        up_index = findfirst(x -> x == name * "up", short_names)
+        if endswith(name, "up")
+            dn_index =
+                findfirst(x -> x == replace(name, "up" => ""), short_names)
+            # Check for updraft vars with no gridmean pair
+            isnothing(dn_index) && push!(grouped_vars, tuple(name))
+        elseif isnothing(up_index)
+            push!(grouped_vars, tuple(name))
         else
-            grouped_vars[info] = [var]
+            push!(grouped_vars, tuple(name, short_names[up_index]))
         end
     end
     return grouped_vars
@@ -740,25 +744,36 @@ 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 = [
+        tuple(
+            collect(
+                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 = [
+        tuple(collect(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,
+        # There is probably a better way to flatten this
+        collect(Iterators.flatten(var_groups_zt)),
         plot_fn = plot_parsed_attribute_title!,
         summary_files = [tmp_file],
         MAX_NUM_COLS = 2,
@@ -772,29 +787,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", "thetaa", "ta", "taup", "haup", "waup", "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 = [
+        tuple(
+            collect(
+                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 = [
+        tuple(collect(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,
+        # There is probably a better way to flatten this
+        collect(Iterators.flatten(var_groups_zt)),
         plot_fn = plot_parsed_attribute_title!,
         summary_files = [tmp_file],
         MAX_NUM_COLS = 2,