Minor fix to resource efficient finetuning plot legends

computational-cell-analytics · Aug 1, 2024 · d0a31eb · d0a31eb
1 parent 0fd30ce
commit d0a31eb
Showing 1 changed file with 16 additions and 10 deletions.
diff --git a/finetuning/specialists/resource-efficient/plot_experiments.py b/finetuning/specialists/resource-efficient/plot_experiments.py
@@ -215,25 +215,31 @@ def _change_opacity_markers(lineplot):
                 ax.get_legend().remove()
 
             _colors = list(PALETTE.values())
-            custom_handles = [
-                mlines.Line2D([], [], color=_colors[0], markersize=15, marker='o', linestyle='-', linewidth=5),
-                mlines.Line2D([], [], color=_colors[0], markersize=15, marker='o', linestyle=':', linewidth=5),
-                mlines.Line2D([], [], color=_colors[1], markersize=15, marker='o', linestyle='-', linewidth=5),
-                mlines.Line2D([], [], color=_colors[1], markersize=15, marker='o', linestyle=':', linewidth=5),
-                mlines.Line2D([], [], color=_colors[2], markersize=15, marker='o', linestyle='-', linewidth=5),
-                mlines.Line2D([], [], color=_colors[2], markersize=15, marker='o', linestyle=':', linewidth=5),
-                mlines.Line2D([], [], color=_colors[3], markersize=15, marker='o', linestyle='-', linewidth=5),
-                mlines.Line2D([], [], color=_colors[3], markersize=15, marker='o', linestyle=':', linewidth=5),
-            ]
 
             if resource_name == "GTX1080":
+                custom_handles = [
+                    mlines.Line2D([], [], color=_colors[0], markersize=15, marker='o', linestyle='-', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[1], markersize=15, marker='o', linestyle='-', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[2], markersize=15, marker='o', linestyle='-', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[3], markersize=15, marker='o', linestyle='-', linewidth=5),
+                ]
                 fig.legend(
                     handles=custom_handles,
                     labels=['AIS (MD, PE)', 'AMG (MD, PE)', 'Point (MD, PE)', 'Box (MD, PE)'],
                     loc="lower center", ncols=4, bbox_to_anchor=(0.5, 0)
                 )
                 bottom = 0.12
             else:
+                custom_handles = [
+                    mlines.Line2D([], [], color=_colors[0], markersize=15, marker='o', linestyle='-', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[0], markersize=15, marker='o', linestyle=':', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[1], markersize=15, marker='o', linestyle='-', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[1], markersize=15, marker='o', linestyle=':', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[2], markersize=15, marker='o', linestyle='-', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[2], markersize=15, marker='o', linestyle=':', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[3], markersize=15, marker='o', linestyle='-', linewidth=5),
+                    mlines.Line2D([], [], color=_colors[3], markersize=15, marker='o', linestyle=':', linewidth=5),
+                ]
                 fig.legend(
                     handles=custom_handles,
                     labels=[