Experiments look good

plotting/plotting.py

@@ -1,4 +1,5 @@
 import seaborn as sns
+import matplotlib as mpl
 from matplotlib import pyplot as plt
 
 PALETTE = 10 * (
@@ -21,7 +22,14 @@ PALETTE = 10 * (
 def plot(filepath, ext='png'):
     plt.tight_layout()
     figure = plt.gcf()
-    figure.savefig(str(filepath), format=ext)
+    ax = plt.gca()
+    legends = [c for c in ax.get_children() if isinstance(c, mpl.legend.Legend)]
+
+    if legends:
+        figure.savefig(str(filepath), format=ext,  bbox_extra_artists=(*legends,), bbox_inches='tight')
+    else:
+        figure.savefig(str(filepath), format=ext)
+
     plt.show()
     plt.clf()
 
@@ -30,7 +38,7 @@ def prepare_tex(df, hue, style, hue_order):
     sns.set(rc={'text.usetex': True}, style='whitegrid')
     lineplot = sns.lineplot(data=df, x='Episode', y='Score', ci=95, palette=PALETTE,
                             hue_order=hue_order, hue=hue, style=style)
-    # lineplot.set_title(f'{sorted(list(df["Measurement"].unique()))}')
+    lineplot.set_title(f'{sorted(list(df["Measurement"].unique()))}')
     plt.legend(bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0)
     plt.tight_layout()
     return lineplot
@@ -48,6 +56,19 @@ def prepare_plt(df, hue, style, hue_order):
     return lineplot
 
 
+def prepare_center_double_column_legend(df, hue, style, hue_order):
+    print('Struggling to plot Figure using LaTeX - going back to normal.')
+    plt.close('all')
+    sns.set(rc={'text.usetex': False}, style='whitegrid')
+    fig = plt.figure(figsize=(10, 11))
+    lineplot = sns.lineplot(data=df, x='Episode', y='Score', hue=hue, style=style,
+                            ci=95, palette=PALETTE, hue_order=hue_order, legend=False)
+    # plt.legend(bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0)
+    lineplot.legend(hue_order, ncol=3, loc='lower center', title='Parameter Combinations', bbox_to_anchor=(0.5, -0.43))
+    plt.tight_layout()
+    return lineplot
+
+
 def prepare_plot(filepath, results_df, ext='png', hue='Measurement', style=None, use_tex=False):
     df = results_df.copy()
     df[hue] = df[hue].str.replace('_', '-')