fixup

src/Visualization/Python/PlotMemoryMonitors.py

@@ -4,6 +4,7 @@
 # See LICENSE.txt for details.
 import logging
 import os
+from math import inf
 from typing import Optional, Sequence
 
 import click
@@ -159,6 +160,9 @@ def compute_individual_total(df: pd.DataFrame):
 
         return total
 
+    # Need .dat because all other components have that extension
+    the_rest_str = "The Rest.dat"
+
     # Given a dictionary of DataFrames, sum all the columns from all the
     # DataFrames that list totals for each node to get a grand total memory
     # usage. Then return a DataFrame with the time and this new total
@@ -172,10 +176,30 @@ def make_totals_df(component_df_dict: dict):
                 totals_df["Time"] = component_df["Time"]
                 totals_df["Total"] = compute_individual_total(component_df)
 
-        return totals_df
+        average_max = totals_df["Total"].mean()
+        components_to_plot = []
+
+        # If the max of a component is too small (less than 1% here) just
+        # group them all together in a category called "The Rest". Otherwise
+        # keep the original component
+        for component_name, component_df in component_df_dict.items():
+            component_total = compute_individual_total(component_df)
+            if component_total.max() < 0.01 * average_max:
+                if the_rest_str in totals_df:
+                    totals_df[the_rest_str] += component_total
+                else:
+                    totals_df[the_rest_str] = component_total.copy()
+            else:
+                totals_df[component_name] = component_total
+                components_to_plot.append(component_name)
+
+        if the_rest_str in totals_df:
+            components_to_plot.append(the_rest_str)
+
+        return totals_df, components_to_plot
 
     # Now make a DataFrame for the total memory across all components
-    totals_df = make_totals_df(df_dict)
+    totals_df, components_to_plot = make_totals_df(df_dict)
 
     # For plotting in MB vs GB
     divisor = 1.0 if use_mb else 1000.0
@@ -187,32 +211,62 @@ def make_totals_df(component_df_dict: dict):
     fig = plt.figure()
     ax = fig.add_subplot(111)
 
-    ax.plot(totals_df["Time"], totals_df["Total"] / divisor)
+    # If we are coloring components, add this to the legend
+    if color_components:
+        ax.plot(
+            totals_df["Time"],
+            totals_df["Total"] / divisor,
+            color="black",
+            label="Total",
+        )
+    else:
+        ax.plot(totals_df["Time"], totals_df["Total"] / divisor, color="black")
 
-    for component in df_dict:
-        # Plot in color (with a legend) for every component if specified
+    for component in components_to_plot:
+        # Plot in color (with a legend)
         if color_components:
             ax.plot(
-                df_dict[component]["Time"],
-                compute_individual_total(df_dict[component]) / divisor,
-                linewidth=0.1,
+                totals_df["Time"],
+                totals_df[component] / divisor,
+                linewidth=0.2,
                 # Remove .dat extension
                 label=component[:-4],
             )
         else:
             ax.plot(
-                df_dict[component]["Time"],
-                compute_individual_total(df_dict[component]) / divisor,
+                totals_df["Time"],
+                totals_df[component] / divisor,
                 color="grey",
-                linewidth=0.1,
+                linewidth=0.2,
             )
 
     gb_or_mb = "MB" if use_mb else "GB"
     plt.title(f"Total Memory Usage ({gb_or_mb})")
     if x_label is not None:
         ax.set_xlabel(x_label)
     if color_components:
-        leg = ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+        # Total is always first
+        maxes = [inf]
+        for component in components_to_plot:
+            maxes.append(totals_df[component].max())
+        # Now that we have all the maxes, we need to sort them and get the
+        # permutation of the sort for the legend so it's in order of largest at
+        # the top, smallest at the bottom
+        maxes_tuple = [(maxes[i], i) for i in range(len(maxes))]
+        maxes_tuple.sort()
+        # Reverse because of the order of plotting
+        maxes_tuple.reverse()
+        _, permutation = zip(*maxes_tuple)
+
+        handles, labels = plt.gca().get_legend_handles_labels()
+        leg = plt.legend(
+            [handles[idx] for idx in permutation],
+            [labels[idx] for idx in permutation],
+            loc="center left",
+            bbox_to_anchor=(1, 0.5),
+        )
+        # The lines in the legend are a bit small because of the plot linewidth,
+        # so make the legend lines a bit bigger
         for line in leg.get_lines():
             line.set_linewidth(1.0)
 

tests/Unit/Visualization/Python/CMakeLists.txt

@@ -53,7 +53,7 @@ if(${BUILD_PYTHON_BINDINGS})
   set_tests_properties(
     "Unit.Visualization.Python.PlotMemoryMonitors"
     PROPERTIES
-    TIMEOUT 40
+    TIMEOUT 10
     )
   set_tests_properties(
     "Unit.Visualization.Python.PlotPowerMonitors"