|
| 1 | +import numpy as np |
| 2 | + |
| 3 | +from matplotlib.widgets import LassoSelector |
| 4 | +from matplotlib.path import Path |
| 5 | + |
| 6 | + |
| 7 | +class SelectFromCollection(object): |
| 8 | + """Select indices from a matplotlib collection using `LassoSelector`. |
| 9 | +
|
| 10 | + Selected indices are saved in the `ind` attribute. This tool highlights |
| 11 | + selected points by fading them out (i.e., reducing their alpha values). |
| 12 | + If your collection has alpha < 1, this tool will permanently alter them. |
| 13 | +
|
| 14 | + Note that this tool selects collection objects based on their *origins* |
| 15 | + (i.e., `offsets`). |
| 16 | +
|
| 17 | + Parameters |
| 18 | + ---------- |
| 19 | + ax : :class:`~matplotlib.axes.Axes` |
| 20 | + Axes to interact with. |
| 21 | +
|
| 22 | + collection : :class:`matplotlib.collections.Collection` subclass |
| 23 | + Collection you want to select from. |
| 24 | +
|
| 25 | + alpha_other : 0 <= float <= 1 |
| 26 | + To highlight a selection, this tool sets all selected points to an |
| 27 | + alpha value of 1 and non-selected points to `alpha_other`. |
| 28 | + """ |
| 29 | + def __init__(self, ax, collection, alpha_other=0.3): |
| 30 | + self.canvas = ax.figure.canvas |
| 31 | + self.collection = collection |
| 32 | + self.alpha_other = alpha_other |
| 33 | + |
| 34 | + self.xys = collection.get_offsets() |
| 35 | + self.Npts = len(self.xys) |
| 36 | + |
| 37 | + # Ensure that we have separate colors for each object |
| 38 | + self.fc = collection.get_facecolors() |
| 39 | + if len(self.fc) == 0: |
| 40 | + raise ValueError('Collection must have a facecolor') |
| 41 | + elif len(self.fc) == 1: |
| 42 | + self.fc = np.tile(self.fc, self.Npts).reshape(self.Npts, -1) |
| 43 | + |
| 44 | + self.lasso = LassoSelector(ax, onselect=self.onselect) |
| 45 | + self.ind = [] |
| 46 | + |
| 47 | + def onselect(self, verts): |
| 48 | + path = Path(verts) |
| 49 | + self.ind = np.nonzero([path.contains_point(xy) for xy in self.xys])[0] |
| 50 | + self.fc[:, -1] = self.alpha_other |
| 51 | + self.fc[self.ind, -1] = 1 |
| 52 | + self.collection.set_facecolors(self.fc) |
| 53 | + self.canvas.draw_idle() |
| 54 | + |
| 55 | + def disconnect(self): |
| 56 | + self.lasso.disconnect_events() |
| 57 | + self.fc[:, -1] = 1 |
| 58 | + self.collection.set_facecolors(self.fc) |
| 59 | + self.canvas.draw_idle() |
| 60 | + |
| 61 | + |
| 62 | +if __name__ == '__main__': |
| 63 | + import matplotlib.pyplot as plt |
| 64 | + |
| 65 | + plt.ion() |
| 66 | + data = np.random.rand(100, 2) |
| 67 | + |
| 68 | + subplot_kw = dict(xlim=(0,1), ylim=(0,1), autoscale_on=False) |
| 69 | + fig, ax = plt.subplots(subplot_kw=subplot_kw) |
| 70 | + |
| 71 | + pts = ax.scatter(data[:, 0], data[:, 1], s=80) |
| 72 | + selector = SelectFromCollection(ax, pts) |
| 73 | + |
| 74 | + plt.draw() |
| 75 | + raw_input('Press any key to accept selected points') |
| 76 | + print "Selected points:" |
| 77 | + print selector.xys[selector.ind] |
| 78 | + selector.disconnect() |
| 79 | + |
| 80 | + # Block end of script so you can check that the lasso is disconnected. |
| 81 | + raw_input('Press any key to quit') |
| 82 | + |
0 commit comments