diff --git a/rl_portfolio_management/callbacks/notebook_plot.py b/rl_portfolio_management/callbacks/notebook_plot.py index 7ce6f35..8c5a74c 100644 --- a/rl_portfolio_management/callbacks/notebook_plot.py +++ b/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