""" Show how to use a lasso to select a set of points and get the indices of the selected points. A callback is used to change the color of the selected points This is currently a proof-of-concept implementation (though it is usable as is). There will be some refinement of the API and the inside polygon detection routine. """ from matplotlib.widgets import Lasso import matplotlib.mlab from matplotlib.nxutils import points_inside_poly from matplotlib.colors import colorConverter from matplotlib.collections import RegularPolyCollection from matplotlib.pyplot import figure, show from numpy import nonzero from numpy.random import rand class Datum: colorin = colorConverter.to_rgba('red') colorout = colorConverter.to_rgba('green') def __init__(self, x, y, include=False): self.x = x self.y = y if include: self.color = self.colorin else: self.color = self.colorout class LassoManager: def __init__(self, ax, data): self.axes = ax self.canvas = ax.figure.canvas self.data = data self.Nxy = len(data) facecolors = [d.color for d in data] self.xys = [(d.x, d.y) for d in data] fig = ax.figure self.collection = RegularPolyCollection( fig.dpi, 6, sizes=(100,), facecolors=facecolors, offsets = self.xys, transOffset = ax.transData) ax.add_collection(self.collection) self.cid = self.canvas.mpl_connect('button_press_event', self.onpress) self.ind = None def callback(self, verts): facecolors = self.collection.get_facecolors() ind = nonzero(points_inside_poly(self.xys, verts))[0] for i in range(self.Nxy): if i in ind: facecolors[i] = Datum.colorin else: facecolors[i] = Datum.colorout self.canvas.draw_idle() self.canvas.widgetlock.release(self.lasso) del self.lasso self.ind = ind def onpress(self, event): if self.canvas.widgetlock.locked(): return if event.inaxes is None: return self.lasso = Lasso(event.inaxes, (event.xdata, event.ydata), self.callback) # acquire a lock on the widget drawing self.canvas.widgetlock(self.lasso) if __name__ == '__main__': data = [Datum(*xy) for xy in rand(100, 2)] fig = figure() ax = fig.add_subplot(111, xlim=(0,1), ylim=(0,1), autoscale_on=False) lman = LassoManager(ax, data) show()