plot_model: use the same colors in stacked plot as in general plot

python/memory_contention/comm_comp.py

@@ -774,6 +774,10 @@ class CommCompGraph:
             for i in range(0, nb_zones, 2):
                 self._ax.axvspan((i*x_span_zone) + 0.5, ((i+1)*x_span_zone)+0.5, facecolor="0.9", zorder=0.1)
 
+    @property
+    def last_used_color(self):
+        return self._current_color
+
     @property
     def nb_curves(self):
         return self._nb_comm_curves + self._nb_comp_curves

python/plot_model.py

@@ -148,6 +148,7 @@ for numa_comp in range(topo.nb_numa_nodes_total):
             [parser.comm_bw_with_comp_results[x]['d1'] for x in real_x_values],
             comm_with_comp_upper_values
         )
+        comm_color = graph.last_used_color
 
         if cli_args.model:
             comm_values_model = combined_model.comm_with_comp(numa_comm, numa_comp)
@@ -189,6 +190,7 @@ for numa_comp in range(topo.nb_numa_nodes_total):
             [parser.comp_with_comm_results[op][x][metric]['min'] for x in real_x_values],
             [parser.comp_with_comm_results[op][x][metric]['max'] for x in real_x_values]
         )
+        comp_color = graph.last_used_color
 
         if cli_args.model:
             comp_values_model = combined_model.comp_with_comm(numa_comm, numa_comp)
@@ -212,18 +214,22 @@ for numa_comp in range(topo.nb_numa_nodes_total):
         graphs[numa_comp].append(graph)
 
         # Stacked plot:
+        colors = iter(plt.rcParams["axes.prop_cycle"].by_key()["color"])
+        next(colors)  # skip the two already used colors
+        next(colors)
         fig = plt.figure(1024, clear=True)
         ax = fig.gca()
         plt.stackplot(
             real_x_values,
             [[parser.comp_with_comm_results[op][x][metric]['avg'] for x in real_x_values], [parser.comm_bw_with_comp_results[x]['med'] for x in real_x_values]],
-            labels=["Comp with comm bw", "Comm with comp bw"],
+            labels=["Computations (with communications)", "Communications (with computations)"],
+            colors=[comp_color, comm_color],
         )
-        plt.plot(real_x_values, [parser.comp_alone_results[op][x][metric]['avg'] for x in real_x_values], label="Comp alone bw")
+        plt.plot(real_x_values, [parser.comp_alone_results[op][x][metric]['avg'] for x in real_x_values], color=next(colors), label="Computations (alone)")
         if cli_args.model:
-            plt.plot(real_x_values, comp_values_model, "+", label="Comp with comm bw model")
-            plt.plot(real_x_values, [sum(x) for x in zip(comp_values_model, comm_values_model)], "+", label="Comm with comp bw model")
-            plt.plot(real_x_values, total_model, "+", label="Total model")
+            plt.plot(real_x_values, comp_values_model, "+", color=next(colors), label="Computations (with communications) - model")
+            plt.plot(real_x_values, [sum(x) for x in zip(comp_values_model, comm_values_model)], "+", color=next(colors), label="Communications (with computations) - model")
+            plt.plot(real_x_values, total_model, "+", color=next(colors), label="Total - model")
         ax.legend()
         ax.grid(axis='y')
         ax.xaxis.set_major_locator(MaxNLocator(integer=True))