From 1ce53eb56ac9ef44c954a2aab632c9f5f8815608 Mon Sep 17 00:00:00 2001
From: Kyle Nelli <knelli@caltech.edu>
Date: Tue, 27 Feb 2024 10:58:35 -0800
Subject: [PATCH] fixup

---
 .../Python/PlotMemoryMonitors.py              | 74 ++++++++++++++++---
 .../Unit/Visualization/Python/CMakeLists.txt  |  2 +-
 2 files changed, 65 insertions(+), 11 deletions(-)

diff --git a/src/Visualization/Python/PlotMemoryMonitors.py b/src/Visualization/Python/PlotMemoryMonitors.py
index 60138fec8bf4c..00636f0be7aa0 100644
--- a/src/Visualization/Python/PlotMemoryMonitors.py
+++ b/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,22 +211,31 @@ 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,
+                totals_df["Time"],
+                totals_df[component] / divisor,
                 linewidth=0.1,
                 # 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,
             )
@@ -212,7 +245,28 @@ def make_totals_df(component_df_dict: dict):
     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)
 
diff --git a/tests/Unit/Visualization/Python/CMakeLists.txt b/tests/Unit/Visualization/Python/CMakeLists.txt
index 6a7981980297d..899c8fbf7a929 100644
--- a/tests/Unit/Visualization/Python/CMakeLists.txt
+++ b/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"