Add new navigate modes in matplotlib

I'm writing a wx/matplotlib application, and I'm having considerable difficulty adding a new tool to the matplotlib NavigationToolbar.

Basically I want to add tools for selection (marquee, lasso, etc) that will toggle the controlled subplots mouse mode. As of yet I have been unable to find any features that will let me do this easily.

I did, however, just discover this function that looked like it would be helpful: http://matplotlib.sourceforge.net/api/axes_api.html?highlight=set_navigate_mode#matplotlib.axes.Axes.set_navigate_mode

Unfortunately, as the warning implies, it doesn't really help me.

Does anybody have a clue of how to do this? Below is a stripped down example showing how far I've gotten. The bookmark icon is used in place of my lasso icon, and I've removed the lasso functionality for brevity.

import wx
from matplotlib.patches import Rectangle
from matplotlib.widgets import Lasso
from matplotlib.figure import Figure
from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg
from matplotlib.backends.backend_wxagg import NavigationToolbar2WxAgg as NavigationToolbar

class ScatterPanel(FigureCanvasWxAgg):
    '''
    Contains the guts for drawing scatter plots.
    '''
    def __init__(self, parent, **kwargs):
        self.figure = Figure()
        FigureCanvasWxAgg.__init__(self, parent, -1, self.figure, **kwargs)
        self.canvas = self.figure.canvas
        self.SetMinSize((100,100))
        self.figure.set_facecolor((1,1,1))
        self.figure.set_edgecolor((1,1,1))
        self.canvas.SetBackgroundColour('white')

        self.subplot = self.figure.add_subplot(111)
        self.navtoolbar = None
        self.lasso = None
        self.redraw()

        self.canvas.mpl_connect('button_press_event', self.on_press)
        self.canvas.mpl_connect('button_release_event', self.on_release)

    def lasso_callback(self, verts):
        pass

    def on_press(self, evt):
        if evt.button == 1:
            if self.canvas.widgetlock.locked(): 
                return
            if evt.inaxes is None: 
                return
            if self.navtoolbar.mode == 'lasso':
                self.lasso = Lasso(evt.inaxes, (evt.xdata, evt.ydata), self.lasso_callback)
                self.canvas.widgetlock(self.lasso)

    def on_release(self, evt):
        # Note: lasso_callback is not called on click without drag so we release
        #   the lock here to handle this case as well.
        if evt.button == 1:
            if self.lasso:
                self.canvas.draw_idle()
                self.canvas.widgetlock.release(self.lasso)
                self.lasso = None
        else:
            self.show_popup_menu((evt.x, self.canvas.GetSize()[1]-evt.y), None)

    def redraw(self):
        self.subplot.clear()
        self.subplot.scatter([1,2,3],[3,1,2])

    def toggle_lasso_tool(self, evt):
        if evt.Checked():
            self.navtoolbar.mode = 'lasso'
            #self.subplot.set_navigate_mode('lasso')
            # Cheat: untoggle the zoom and pan tools
            self.navtoolbar.ToggleTool(self.navtoolbar._NTB2_ZOOM, False)
            self.navtoolbar.ToggleTool(self.navtoolbar._NTB2_PAN, False)
        else: 
      开发者_JAVA百科      self.navtoolbar.mode = ''
            self.lasso = None
            #self.subplot.set_navigate_mode('')

    def get_toolbar(self):
        if not self.navtoolbar:
            self.navtoolbar = NavigationToolbar(self.canvas)
            self.navtoolbar.DeleteToolByPos(6)
            ID_LASSO_TOOL = wx.NewId()
            lasso = self.navtoolbar.InsertSimpleTool(5, ID_LASSO_TOOL, 
                            wx.ArtProvider.GetBitmap(wx.ART_ADD_BOOKMARK),
                            isToggle=True)
            self.navtoolbar.Realize()
            self.navtoolbar.Bind(wx.EVT_TOOL, self.toggle_lasso_tool, id=ID_LASSO_TOOL)
        return self.navtoolbar

if __name__ == "__main__":
    app = wx.PySimpleApp()
    f = wx.Frame(None, size=(600,600))
    p = ScatterPanel(f)
    f.SetToolBar(p.get_toolbar())            
    f.Show()
    app.MainLoop()

Thanks, Adam

Here's an improved version of MyNavToolbar. The main thing to note is the addition of the add_user_tool method. I call it from within __init__, but you'd probably want to call it from outside the MyNavToolbar class. This way you could have different tools for plot types.

class MyNavToolbar(NavigationToolbar2WxAgg):
    """wx/mpl NavToolbar hack with an additional tools user interaction.
    This class is necessary because simply adding a new togglable tool to the
    toolbar won't (1) radio-toggle between the new tool and the pan/zoom tools.
    (2) disable the pan/zoom tool modes in the associated subplot(s).
    """
    def __init__(self, canvas):
        super(NavigationToolbar2WxAgg, self).__init__(canvas)
        self.pan_tool  = self.FindById(self._NTB2_PAN)
        self.zoom_tool = self.FindById(self._NTB2_ZOOM)
        self.Bind(wx.EVT_TOOL, self.on_toggle_pan_zoom, self.zoom_tool)
        self.Bind(wx.EVT_TOOL, self.on_toggle_pan_zoom, self.pan_tool)

        self.user_tools = {}   # user_tools['tool_mode'] : wx.ToolBarToolBase

        self.InsertSeparator(5)
        self.add_user_tool('lasso', 6, icons.lasso_tool.ConvertToBitmap(), True, 'Lasso')
        self.add_user_tool('gate', 7, icons.gate_tool.ConvertToBitmap(), True, 'Gate')

    def add_user_tool(self, mode, pos, bmp, istoggle=True, shortHelp=''):
        """Adds a new user-defined tool to the toolbar.
        mode -- the value that MyNavToolbar.get_mode() will return if this tool 
                is toggled on
        pos -- the position in the toolbar to add the icon
        bmp -- a wx.Bitmap of the icon to use in the toolbar
        isToggle -- whether or not the new tool toggles on/off with the other 
                    togglable tools
        shortHelp -- the tooltip shown to the user for the new tool
        """
        tool_id = wx.NewId()
        self.user_tools[mode] = self.InsertSimpleTool(pos, tool_id, bmp,
                            isToggle=istoggle, shortHelpString=shortHelp)
        self.Bind(wx.EVT_TOOL, self.on_toggle_user_tool, self.user_tools[mode])

    def get_mode(self):
        """Use this rather than navtoolbar.mode
        """
        for mode, tool in self.user_tools.items():
            if tool.IsToggled():
                return mode
        return self.mode

    def untoggle_mpl_tools(self):
        """Hack city: Since I can't figure out how to change the way the 
        associated subplot(s) handles mouse events: I generate events to turn
        off whichever tool mode is enabled (if any). 
        This function needs to be called whenever any user-defined tool 
        (eg: lasso) is clicked.
        """
        if self.pan_tool.IsToggled():
            wx.PostEvent(
                self.GetEventHandler(), 
                wx.CommandEvent(wx.EVT_TOOL.typeId, self._NTB2_PAN)
            )
            self.ToggleTool(self._NTB2_PAN, False)
        elif self.zoom_tool.IsToggled():
            wx.PostEvent(
                self.GetEventHandler(),
                wx.CommandEvent(wx.EVT_TOOL.typeId, self._NTB2_ZOOM)
            )
            self.ToggleTool(self._NTB2_ZOOM, False)

    def on_toggle_user_tool(self, evt):
        """user tool click handler.
        """
        if evt.Checked():
            self.untoggle_mpl_tools()
            #untoggle other user tools
            for tool in self.user_tools.values():
                if tool.Id != evt.Id:
                    self.ToggleTool(tool.Id, False)

    def on_toggle_pan_zoom(self, evt):
        """Called when pan or zoom is toggled. 
        We need to manually untoggle user-defined tools.
        """
        if evt.Checked():
            for tool in self.user_tools.values():
                self.ToggleTool(tool.Id, False)
        # Make sure the regular pan/zoom handlers get the event
        evt.Skip()

    def reset_history(self):
        """More hacky junk to clear/reset the toolbar history.
        """
        self._views.clear()
        self._positions.clear()
        self.push_current()

Well here it is, ugly but functional. I'll let the docstrings do the talking, this has wasted enough of my time.

import wx
from matplotlib.patches import Rectangle
from matplotlib.widgets import Lasso
from matplotlib.figure import Figure
from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg
from matplotlib.backends.backend_wxagg import NavigationToolbar2WxAgg

class MyNavToolbar(NavigationToolbar2WxAgg):
    """wx/mpl NavToolbar hack with an additional tools user interaction.
    This class is necessary because simply adding a new togglable tool to the
    toolbar won't (1) radio-toggle between the new tool and the pan/zoom tools.
    (2) disable the pan/zoom tool modes in the associated subplot(s).
    """
    ID_LASSO_TOOL = wx.NewId()
    def __init__(self, canvas):
        super(NavigationToolbar2WxAgg, self).__init__(canvas)

        self.pan_tool  = self.FindById(self._NTB2_PAN)
        self.zoom_tool = self.FindById(self._NTB2_ZOOM)

        self.lasso_tool = self.InsertSimpleTool(5, self.ID_LASSO_TOOL, 
                            wx.ArtProvider.GetBitmap(wx.ART_ADD_BOOKMARK),
                            isToggle=True)
        self.Bind(wx.EVT_TOOL, self.on_toggle_lasso_tool, self.lasso_tool)
        self.Bind(wx.EVT_TOOL, self.on_toggle_pan_zoom, self.zoom_tool)
        self.Bind(wx.EVT_TOOL, self.on_toggle_pan_zoom, self.pan_tool)

    def get_mode(self):
        """Use this rather than navtoolbar.mode
        """
        if self.lasso_tool.IsToggled():
            return 'lasso'
        else:
            return self.mode

    def untoggle_mpl_tools(self):
        """Hack city: Since I can't figure out how to change the way the 
        associated subplot(s) handles mouse events: I generate events to turn
        off whichever tool mode is enabled (if any). 
        This function needs to be called whenever any user-defined tool 
        (eg: lasso) is clicked.
        """
        if self.pan_tool.IsToggled():
            wx.PostEvent(
                self.GetEventHandler(), 
                wx.CommandEvent(wx.EVT_TOOL.typeId, self._NTB2_PAN)
            )
            self.ToggleTool(self._NTB2_PAN, False)
        elif self.zoom_tool.IsToggled():
            wx.PostEvent(
                self.GetEventHandler(),
                wx.CommandEvent(wx.EVT_TOOL.typeId, self._NTB2_ZOOM)
            )
            self.ToggleTool(self._NTB2_ZOOM, False)

    def on_toggle_lasso_tool(self, evt):
        """Lasso tool handler.
        """
        if evt.Checked():
            self.untoggle_mpl_tools()

    def on_toggle_pan_zoom(self, evt):
        """Called when pan or zoom is toggled. 
        We need to manually untoggle user-defined tools.
        """
        if evt.Checked():
            self.ToggleTool(self.ID_LASSO_TOOL, False)
        # Make sure the regular pan/zoom handlers get the event
        evt.Skip()

class ScatterPanel(FigureCanvasWxAgg):
    """Contains the guts for drawing scatter plots.
    """
    def __init__(self, parent, **kwargs):
        self.figure = Figure()
        FigureCanvasWxAgg.__init__(self, parent, -1, self.figure, **kwargs)
        self.canvas = self.figure.canvas
        self.SetMinSize((100,100))
        self.figure.set_facecolor((1,1,1))
        self.figure.set_edgecolor((1,1,1))
        self.canvas.SetBackgroundColour('white')

        self.subplot = self.figure.add_subplot(111)
        self.navtoolbar = None
        self.lasso = None
        self.redraw()

        self.canvas.mpl_connect('button_press_event', self.on_press)
        self.canvas.mpl_connect('button_release_event', self.on_release)

    def lasso_callback(self, verts):
        pass

    def on_press(self, evt):
        """canvas mousedown handler
        """
        if evt.button == 1:
            if self.canvas.widgetlock.locked(): 
                return
            if evt.inaxes is None: 
                return
            if self.navtoolbar and self.navtoolbar.get_mode() == 'lasso':
                self.lasso = Lasso(evt.inaxes, (evt.xdata, evt.ydata), self.lasso_callback)
                self.canvas.widgetlock(self.lasso)

    def on_release(self, evt):
        """canvas mouseup handler
        """
        # Note: lasso_callback is not called on click without drag so we release
        #   the lock here to handle this case as well.
        if evt.button == 1:
            if self.lasso:
                self.canvas.draw_idle()
                self.canvas.widgetlock.release(self.lasso)
                self.lasso = None
        else:
            self.show_popup_menu((evt.x, self.canvas.GetSize()[1]-evt.y), None)

    def redraw(self):
        self.subplot.clear()
        self.subplot.scatter([1,2,3],[3,1,2])

    def get_toolbar(self):
        if not self.navtoolbar:
            self.navtoolbar = MyNavToolbar(self.canvas)
            self.navtoolbar.Realize()
        return self.navtoolbar

if __name__ == "__main__":
    app = wx.PySimpleApp()
    f = wx.Frame(None, size=(600,600))
    p = ScatterPanel(f)
    f.SetToolBar(p.get_toolbar())            
    f.Show()
    app.MainLoop()