Make order of events more intuitive in profiling timeline plot

multiscale · Oct 1, 2023 · 039e8d6 · 039e8d6
1 parent bfb8bd7
commit 039e8d6
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/muscle3/profiling.py b/muscle3/profiling.py
@@ -184,7 +184,7 @@ def __init__(self, performance_file: Path) -> None:
         ax.set_xlabel('Wallclock time (s)')
 
         # Background
-        ax.barh(
+        running_artist = ax.barh(
                 instances,
                 [end_times[i] - begin_times[i] for i in instances],
                 _BAR_WIDTH,
@@ -230,7 +230,13 @@ def __init__(self, performance_file: Path) -> None:
         ax.set_autoscale_on(True)
         ax.callbacks.connect('xlim_changed', self.update_data)
 
-        ax.legend(loc='upper right')
+        ordered_artists = [self._bars[event_type] for event_type in _EVENT_TYPES]
+        ordered_names = list(_EVENT_TYPES)
+
+        ordered_artists.insert(6, running_artist)
+        ordered_names.insert(6, 'RUNNING')
+
+        ax.legend(ordered_artists, ordered_names, loc='upper right')
         ax.figure.canvas.draw_idle()
 
     def close(self) -> None: