move legend

rl_portfolio_management/callbacks/notebook_plot.py

@@ -17,7 +17,7 @@ class LivePlotNotebook(object):
     liveplot.update(x, [ya,yb])
     """
 
-    def __init__(self, log_dir=None, episode=0, labels=[], title='', ylabel='returns', colors=None, linestyles=None):
+    def __init__(self, log_dir=None, episode=0, labels=[], title='', ylabel='returns', colors=None, linestyles=None, legend_outside=True):
         if not matplotlib.rcParams['backend'] == 'nbAgg':
             logging.warn("The liveplot callback only work when matplotlib is using the nbAgg backend. Execute 'matplotlib.use('nbAgg', force=True)'' or '%matplotlib notebook'")
 
@@ -43,12 +43,18 @@ class LivePlotNotebook(object):
 
         ax.set_xlim(0, 1)
         ax.set_ylim(0, 1)
-        ax.legend()
         ax.set_xlabel('date')
         ax.set_ylabel(ylabel)
         ax.grid()
         ax.set_title(title)
 
+        # give the legend it's own space, the right 20% where it right align left
+        if legend_outside:
+            fig.subplots_adjust(right=0.8)
+            ax.legend(loc='center left', bbox_to_anchor=(1.0, 0.5), frameon=False)
+        else:
+            ax.legend()
+
         self.ax = ax
         self.fig = fig