Dynamic of errorbar plot in matplotlib

#!/usr/bin/env python

"""A wxPython widget to dislay a simple matplotlib graph.

This could be made *much* more complicated.

TODO:
    Allow series colour change
    Allow line, points, etc, series
    Allow changing of number of points in series (how?)
    Allow more than one series (how?)
    Allow setting of axis ticks/subticks
    etc, etc

    The basic idea is to minimise the number of args passed to the constructor
    and add methods to control the graph look&feel.
"""


import wx
import numpy as np
from matplotlib.figure import Figure
from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas
from matplotlib.backends.backend_wx import NavigationToolbar2Wx
from matplotlib.path import Path
import matplotlib.patches as patches
import pylab


class MplWidget(wx.Panel):
    """Class to display a matplotlib graph on a wxPython panel.

    Note that the dataset length must be set initially AND NOT CHANGED
    THEREAFTER.  Whether this is fixed in stone or merely due to my
    imperfect understanding of wxPython ...
    """

    def __init__(self, parent, toolbar=False, num_points=100,
                 x_range=(-1,1), y_range=(-1,1), **kwargs):
        """Create an MplWidget instance.

        parent      the owning object reference
        toolbar     if True attach the matplotlib toolbar to graph
        num_points  number of points in the dataset
        x_range     tuple (min, max) of graph X values
        y_range     tuple (min, max) of graph Y values
        kwargs      panel-specific keyword params

        We set the data ranges once instead of letting each update automatically
        set them - this can lead to strangeness in axis labelling.  Consider
        adding a self.setRanges() function if you need to change the X/Y ranges
        for an update.
        """

        wx.Panel.__init__(self, parent, **kwargs)

        # create a static vertical boxsizer, a nice enclosing line
        self.sbsizer = wx.StaticBoxSizer(wx.StaticBox(self, wx.ID_ANY, ''),
                                         orient=wx.VERTICAL)

        # usual Matplotlib functions
        self.figure = Figure(figsize=(2,2))
        self.axes = self.figure.add_subplot(111)
        self.patch = None

        # define the 2D dataset for number of points and containing None
        # it appears we can't change a dataset *size* dynamically!?
        self.x_data = [None] * num_points
        self.y_data = [None] * num_points
        self.y_error = [None] * num_points
        self.plot_data = self.axes.plot(self.x_data, self.y_data,
                                        color='red', lw=3)[0]

        # graph title and labels - set to nondescript values
        # user can change these
        pylab.setp(self.axes, title='')
        pylab.setp(self.axes, xlabel='')
        pylab.setp(self.axes, ylabel='')

        # set X and Y limits
        # these could be changed, but we don't allow it in this version
        (x_min, x_max) = x_range
        self.axes.set_xbound(lower=x_min, upper=x_max)
        (y_min, y_max) = y_range
        self.axes.set_ybound(lower=y_min, upper=y_max)

        # initialize the FigureCanvas, mapping the figure to the WxAgg backend
        self.canvas = FigureCanvas(self, wx.ID_ANY, self.figure)

        # add the figure canvas
        self.sbsizer.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.EXPAND)

        # if the user wants a toolbar
        if toolbar:
            # instantiate the Navigation Toolbar
            self.toolbar = NavigationToolbar2Wx(self.canvas)
            # needed to support Windows systems
            self.toolbar.Realize()

            # add it to the sub-box sizer
            self.sbsizer.Add(self.toolbar, 0, wx.LEFT)
            # explicitly show the toolbar
            self.toolbar.Show()

        self.SetSizer(self.sbsizer)


    def setData(self, x_data, y_data, e_data):
        """Update the widget X, Y and errorbar data."""

        # change graph data
        self.x_data = x_data
        self.y_data = y_data
        self.plot_data.set_data(np.array(x_data), np.array(y_data))

        # remove any previous errorbar patch
        if self.patch:
            self.axes.patches.remove(self.patch)

        # calculate new errorbar info, patch it in
        verts = []
        codes = []
        for (x, y, e) in zip(x_data, y_data, e_data):
            verts.append((x, y+e))
            codes.append(Path.MOVETO)

            verts.append((x, y-e))
            codes.append(Path.LINETO)
        barpath = Path(verts, codes)
        self.patch = patches.PathPatch(barpath, facecolor='grey', alpha=0.5)
        self.axes.add_patch(self.patch)


    def setTitle(self, title):
        """Set a new graph title."""

        pylab.setp(self.axes, title=title)


    def setAxisLabels(self, x_label, y_label=None):
        """Set new graph axis labels."""

        pylab.setp(self.axes, xlabel=x_label, ylabel=y_label)


    def setGrid(self, state):
        """Fiddle with the graph grid.

        state   if True show the grid
        """

        self.axes.grid(state)


    def setLinewidth(self, linewidth=1):
        """Set line series linewidth."""

        pylab.setp(self.plot_data, linewidth=linewidth)


    def Refresh(self):
        """Draw the series plot data specified.

        We don't do a refresh automatically after each set*() change.
        That could lead to many unnecessary redraws and flashing.
        Let the user control do:
            possibly many updates, then
            one refresh
        """

        self.canvas.draw()


if __name__ == '__main__':
    class MyForm(wx.Frame):
        def __init__(self):
            wx.Frame.__init__(self, None, wx.ID_ANY, title='Errorbar Test')
            self.SetMinSize((550,450))

            self.panel = wx.Panel(self, wx.ID_ANY)

            topSizer = wx.BoxSizer(wx.VERTICAL)
            self.mpl = MplWidget(self.panel, num_points=5,
                                 x_range=(0,4), y_range=(0,1))
            self.mpl.setTitle('Test graph')
            self.mpl.setAxisLabels('X-axis', 'Y-axis')
            self.mpl.setGrid(True)
            topSizer.Add(self.mpl, 1, wx.ALL|wx.EXPAND)

            self.btn = wx.Button(self.panel, -1, 'Another')
            topSizer.Add(self.btn, 0, wx.ALIGN_RIGHT)

            self.panel.SetSizer(topSizer)
            topSizer.Fit(self)

            self.linewidth = 0

            self.btn.Bind(wx.EVT_BUTTON, self.onAnother)

        def onAnother(self, event=None):
            x = np.arange(0.0, 5, 1.0)
            y = np.exp(-x) + np.random.random(x.size) / 5
            e = np.random.random(x.size) / 5
            self.mpl.setData(x, y, e)
            self.linewidth += 1
            if self.linewidth > 5:
                self.linewidth = 1
            self.mpl.setTitle('Test graph - linewidth %d' % self.linewidth)
            self.mpl.setLinewidth(self.linewidth)
            self.mpl.Refresh()

    app = wx.PySimpleApp()
    frame = MyForm().Show()
    app.MainLoop()