从散点图构建数组

时间:2014-05-14 12:52:18

标签: python graph matplotlib

我想实现一个程序/函数,我可以从散点图中选择数据,然后将它们存储到一个数组中。关于如何开始或在哪里找到关于该主题的更多信息的任何提示都会很好:)。

1 个答案:

答案 0 :(得分:0)

通过套索或矩形选择进行选择。应该相当容易扩展。仅使用matplotlib小部件。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button, LassoSelector, RectangleSelector
from matplotlib.path import Path
from matplotlib.patches import Rectangle
from functools import partial

class LassoSelect(object):

    def __init__(self, ax, collection):
        self.canvas = ax.figure.canvas
        self.collection = collection
        self.xys = collection.get_offsets()
        self.lasso = LassoSelector(ax, onselect=self.onselect)
        self.setActive(False)

    def onselect(self, verts):
        path = Path(verts)
        ind = np.nonzero([path.contains_point(xy) for xy in self.xys])[0]
        fc = self.collection.get_facecolors()
        fc[ind, -1] = 1
        self.collection.set_facecolors(fc)
        self.canvas.draw_idle()

    def setActive(self, activate):
        if activate:
            self.lasso.active = True
        else:
            self.lasso.active = False


class RectangleSelect(object):

    def __init__(self, ax, collection):
        super(RectangleSelect, self).__init__()
        self.RS = RectangleSelector(ax, self.onselect,
                                    drawtype='box', useblit=True,
                                    button=[1, 3],  # don't use middle button
                                    minspanx=5, minspany=5,
                                    spancoords='pixels')

        self.canvas = ax.figure.canvas
        self.collection = collection
        self.xys = collection.get_offsets()
        self.setActive(False)

    def onselect(self, eclick, erelease):
        'eclick and erelease are the press and release events'
        x1, y1 = eclick.xdata, eclick.ydata
        x2, y2 = erelease.xdata, erelease.ydata
        x0 = min(x1, x2)
        y0 = min(y1, y2)
        width = abs(x2-x1)
        height = abs(y2-y1)
        rect = Rectangle((x0, y0), width, height)

        ind = np.nonzero([rect.contains_point(xy) for xy in self.xys])[0]
        fc = self.collection.get_facecolors()
        fc[ind, -1] = 1
        self.collection.set_facecolors(fc)
        self.canvas.draw_idle()

    def setActive(self, activate):
        if activate:
            self.RS.set_active(True)
        else:
            self.RS.set_active(False)


class DataSelector(object):

    activeColor = (0, 1, 0)
    inActiveColor = (1, 1, 1)

    def __init__(self, x, y):
        super(DataSelector, self).__init__()
        self.x, self.y = x, y
        self.fig, self.ax = plt.subplots()
        self.pts = self.ax.scatter(x, y)
        self.selected_pts = []
        self.addButtons()
        self.selectorWidgets = {'Rectangle': RectangleSelect(self.ax, self.pts),
                                'Lasso': LassoSelect(self.ax, self.pts)}

        self.fc = self.pts.get_facecolors()
        if len(self.fc) == 0:
            raise ValueError('Collection must have a facecolor')
        elif len(self.fc) == 1:
            self.fc = np.tile(self.fc, len(x)).reshape(len(x), -1)
        self.fc[:, -1] = 0.3
        self.pts.set_facecolors(self.fc)

    def resetSelection(self, event=None):
        self.fc[:, -1] = 0.3
        self.fig.canvas.draw()

    def exportSelection(self, event=None):
        ind = np.nonzero(self.fc[:, -1] != 0.3)[0]
        print self.x[ind], self.y[ind]

    def activateWidget(self, key, event):
        for widgetKey, widget in self.selectorWidgets.iteritems():
            if widgetKey == key:
                widget.setActive(True)
            else:
                widget.setActive(False)
        for ax, button in self.buttonAxes.iteritems():
            if ax == event.inaxes:
                button.color = self.activeColor
            else:
                button.color = self.inActiveColor
            button.hovercolor = button.color

    def addButtons(self):
        fig = self.fig
        fig.subplots_adjust(bottom=0.25)

        #Button text and callback
        buttons = [("Lasso", None), ("Rectangle", None), ('Reset', self.resetSelection), ('Export Selection', self.exportSelection)]

        nButtons = len(buttons)
        axBorders = np.linspace(0.25, 0.90, nButtons+1, endpoint=False)
        spacing = axBorders[1] - axBorders[0]
        self.buttonAxes = {}
        for i, btn in zip(axBorders, buttons):
            #[left, bottom, width, height]
            buttonAx = plt.axes([i, 0.1, spacing, 0.04])
            button = Button(buttonAx, btn[0], color=self.inActiveColor)
            if btn[1] is None:
                button.on_clicked(partial(self.activateWidget, btn[0]))
            else:
                button.on_clicked(btn[1])
            self.buttonAxes[buttonAx] = button

x, y = np.random.random((2, 25))*10

dataSelector = DataSelector(x, y)
plt.show()