LivePlot

# Class for doing live plots in IPython

class LivePlot():
    def __init__(self, max_width=None):
        self.x = [0]
        self.y = {}
        self.lines, self.fig, self.axes = None, None, None
        self.max_width = max_width

    def update(self, losses):
        if self.lines is None:
            n = len(losses)
            self.fig, self.axes = plt.subplots(n, 1)
            self.fig.tight_layout()
            # handle case where len(update)==1
            if not isinstance(self.axes, np.ndarray):
                self.axes = np.array([self.axes])
            for ax in self.axes:
                ax.plot([], [])
            self.lines = {k: ax.lines[0] for k, ax in zip(losses.keys(), self.axes)}
            self.axes = {key: ax for ax, key in zip(self.axes, losses.keys())}
            for key, ax in self.axes.items():
                ax.set_title(key)
            self.y = {k: [v] for k, v in losses.items()}
        # update
        else:
            self.x = list(range(len(self.x)+1))
            for k, v in losses.items():
                y, line, ax = self.y[k], self.lines[k], self.axes[k]
                y.append(v)
                line.set_xdata(self.x)
                line.set_ydata(self.y[k])
                if self.max_width is None:
                    ax.set_xlim(min(self.x), max(self.x))
                    ax.set_ylim(min(y), max(y))
                else:
                    x = self.x[-self.max_width:]
                    y_ = y[-self.max_width:]
                    ax.set_xlim(min(x), max(x))
                    ax.set_ylim(min(y_), max(y_))
            self.fig.canvas.draw()