Interactively select points from a plot in matplotlib

In Matlab, there is a very useful function called ginput(). There is no such ready-built function for matplotlib/pylab, but all the pieces are there. Copy and paste this code to get ginput functionality.

The code is below if you want to jump right to it.

Here’s a little explanation of what’s happening:

The general idea is to create a class, the MouseMonitor, that will hold the x and y data coordinates. This class has a single method, the mycall() method. Note: you could be tricky and define a __call__ method, and save a few keystrokes by just passing mouse (instead of mouse.mycall) to the connect function. I didn’t do that here cause it’s harder to explain.

The connect() function needs, as its first argument, the name of an event. Possible events are ‘resize_event’, ‘draw_event’, ‘key_press_event’, ‘key_release_event’, ‘button_press_event’, ‘button_release_event’, ‘motion_notify_event’, and ‘pick_event’.

The second argument to connect() needs to be the name of a function that will be called every time that event happens. In this case, it’s the mycall method of the mouse object, mouse.mycall.

Since mouse.mycall is the callback function given to connect, it will be called every time there is a button press event. So every time you click on the plot, mouse.mycall it will append the x and y coords to its own internal list, print the x and y coords, plot a red circle, and refresh the plot every time.

When you’re done, you can access the x and y data in mouse with mouse.xdatalist and mouse.ydatalist.

More examples to follow, but for now here’s a basic version that improve a little on the functionality of Matlab’s ginput (which doesn’t give you graphical feedback on its own).

from pylab import *
class MouseMonitor:
    event = None
    xdatalist = []
    ydatalist = []

    def mycall(self, event):
        self.event = event
        self.xdatalist.append(event.xdata)
        self.ydatalist.append(event.ydata)

        print 'x = %s and y = %s' % (event.xdata,event.ydata)

        ax = gca()  # get current axis
        ax.hold(True) # overlay plots.

        # Plot a red circle where you clicked.
        ax.plot([event.xdata],[event.ydata],'ro')

        draw()  # to refresh the plot.

if __name__=="__main__":
    # Example usage
    mouse = MouseMonitor()
    connect('button_press_event', mouse.mycall)
    plot([1,2,3])
    show()

An improved version

The following code allows you to set up a MouseMonitor object that plots multiple sets of data.

from pylab import *
import sys
class MouseMonitor:
    event = None
    xdatalist = []
    ydatalist = []

    def my_call(self, event):
        self.event = event

        if event.button == 1:
            self.xdatalist.append(event.xdata)
            self.ydatalist.append(event.ydata)

            sys.stdout.flush()
            print 'x = %s and y = %s' % (event.xdata,event.ydata)
            sys.stdout.flush()
            ax = gca()
            ax.hold(True)
            ax.plot([event.xdata],[event.ydata],'ro')

            draw()

        if event.button == 3:
            disconnect(self.cid)
            self.ask_for_another()

    def ask_for_another(self):
        sys.stdout.flush()
        answer = raw_input('Press n for the next plot, or any other key to stop:  ')
        if answer == 'n':
            self.plot_next()
        else:
            print '::sniff::  bye'
            close('all')

    def set_data(self,x,y):
        self.xlist = x
        self.ylist = y
        self.i = 0

    def plot_next(self):
        if self.i < len(self.xlist):
            cla()
            plot(xlist[self.i],ylist[self.i])
            self.i += 1
            self.connect_to_plot()
            draw()
        else:
            sys.stdout.flush()
            print 'all out of data.'

    def connect_to_plot(self):
        self.cid = connect('button_press_event',self.my_call)

if __name__ == "__main__":
    mouse = MouseMonitor()

    xlist = [[1,2,3],[4,5,6]]
    ylist = [[3,4,5],[8,9,7]]

    mouse.set_data(xlist,ylist)

    mouse.plot_next()

    show()