Change colour of data points on selection and remove them with key press in matplotlib 3d scatter plot

Question

I have a 3d scatter plot in matplotlib, and have set up annotations, inspired by answers here, particularly that by Don Cristobal.

I have some basic event-capturing code set up, but after several days of trying, I still have not managed to achieve my objectives. These are:

(i) Change colour of a point (dot) when selected with left mouse button from blue to e.g. dark blue/green.

(ii) Remove any selected dot selected in (i) after pressing the 'delete' key, including any annotations

(iii) Select multiple points in (i) using a selection rectangle and delete using 'delete' key

I have tried many approaches, including animating the chart to update based on changes in data, manipulating artist parameters, changing data points via e.g. xs, ys, zs = graph._offsets3d (which does not appear to be documented), but to no avail.

I have attempted, within the onpick(event) function, to:

(i) Interact with points via event.ind to change colour using event.artist.set_face_colour()

(ii) Remove points using both artist.remove()

(iii) Remove points using xs, ys, zs = graph._offsets3d , removing the relevant point by index (event.ind[0]) from xs, ys, and zs, and then resetting graph points via graph._offsets3d = xs_new, ys_new, zs_new

(iv) Redrawing the chart, or relevant sections of the chart only (blitting?)

with no success!

My current code is roughly as below. In fact, I have several hundred points, not the 3 in the simplified example below. I would like the graph to update smoothly if possible, although just getting something usable would be great. Most of the code to do this should probably reside within 'onpick', as that is the function that deals with picking events (see event handler). I have retained some of my code attempts, commented out, which I hope may be of some use. The 'forceUpdate' function is meant to update the graph object on an event trigger, but I am not convinced that it currently does anything. function on_key(event) also does not currently seem to work: presumably there must be a setting in order to determine points to delete, e.g. all artists that have a facecolor that has been changed from the default (e.g. delete all points that have colour dark blue/green rather than light blue).

Any help is much appreciated.

The code (below) is called with:

visualize3DData (Y, ids, subindustry)

Some sample data points are below:

#Datapoints
Y = np.array([[ 4.82250000e+01,  1.20276889e-03,  9.14501289e-01], [ 6.17564688e+01,  5.95020883e-02, -1.56770827e+00], [ 4.55139000e+01,  9.13454423e-02, -8.12277299e+00]])

#Annotations
ids = ['a', 'b', 'c']

subindustry =  'example'

My current code is here:

import matplotlib.pyplot as plt, numpy as np
from mpl_toolkits.mplot3d import proj3d

def visualize3DData (X, ids, subindus):
    """Visualize data in 3d plot with popover next to mouse position.

    Args:
        X (np.array) - array of points, of shape (numPoints, 3)
    Returns:
        None
    """
    fig = plt.figure(figsize = (16,10))
    ax = fig.add_subplot(111, projection = '3d')
    graph  = ax.scatter(X[:, 0], X[:, 1], X[:, 2], depthshade = False, picker = True)  

    def distance(point, event):
        """Return distance between mouse position and given data point

        Args:
            point (np.array): np.array of shape (3,), with x,y,z in data coords
            event (MouseEvent): mouse event (which contains mouse position in .x and .xdata)
        Returns:
            distance (np.float64): distance (in screen coords) between mouse pos and data point
        """
        assert point.shape == (3,), "distance: point.shape is wrong: %s, must be (3,)" % point.shape

        # Project 3d data space to 2d data space
        x2, y2, _ = proj3d.proj_transform(point[0], point[1], point[2], plt.gca().get_proj())
        # Convert 2d data space to 2d screen space
        x3, y3 = ax.transData.transform((x2, y2))

        return np.sqrt ((x3 - event.x)**2 + (y3 - event.y)**2)


    def calcClosestDatapoint(X, event):
        """"Calculate which data point is closest to the mouse position.

        Args:
            X (np.array) - array of points, of shape (numPoints, 3)
            event (MouseEvent) - mouse event (containing mouse position)
        Returns:
            smallestIndex (int) - the index (into the array of points X) of the element closest to the mouse position
        """
        distances = [distance (X[i, 0:3], event) for i in range(X.shape[0])]
        return np.argmin(distances)


    def annotatePlot(X, index, ids):
        """Create popover label in 3d chart

        Args:
            X (np.array) - array of points, of shape (numPoints, 3)
            index (int) - index (into points array X) of item which should be printed
        Returns:
            None
        """
        # If we have previously displayed another label, remove it first
        if hasattr(annotatePlot, 'label'):
            annotatePlot.label.remove()
        # Get data point from array of points X, at position index
        x2, y2, _ = proj3d.proj_transform(X[index, 0], X[index, 1], X[index, 2], ax.get_proj())
        annotatePlot.label = plt.annotate( ids[index],
            xy = (x2, y2), xytext = (-20, 20), textcoords = 'offset points', ha = 'right', va = 'bottom',
            bbox = dict(boxstyle = 'round,pad=0.5', fc = 'yellow', alpha = 0.5),
            arrowprops = dict(arrowstyle = '->', connectionstyle = 'arc3,rad=0'))
        fig.canvas.draw()


    def onMouseMotion(event):
        """Event that is triggered when mouse is moved. Shows text annotation over data point closest to mouse."""
        closestIndex = calcClosestDatapoint(X, event)
        annotatePlot (X, closestIndex, ids) 


    def onclick(event):
        print('%s click: button=%d, x=%d, y=%d, xdata=%f, ydata=%f' %
              ('double' if event.dblclick else 'single', event.button,
               event.x, event.y, event.xdata, event.ydata))

    def on_key(event):
        """
        Function to be bound to the key press event
        If the key pressed is delete and there is a picked object,
        remove that object from the canvas
        """
        if event.key == u'delete':
            ax = plt.gca()
            if ax.picked_object:
                ax.picked_object.remove()
                ax.picked_object = None
                ax.figure.canvas.draw()

    def onpick(event):

        xmouse, ymouse = event.mouseevent.xdata, event.mouseevent.ydata
        artist = event.artist
        # print(dir(event.mouseevent))
        ind = event.ind
        # print('Artist picked:', event.artist)
        # # print('{} vertices picked'.format(len(ind)))
        print('ind', ind)
        # # print('Pick between vertices {} and {}'.format(min(ind), max(ind) + 1))
        # print('x, y of mouse: {:.2f},{:.2f}'.format(xmouse, ymouse))
        # # print('Data point:', x[ind[0]], y[ind[0]])
        #
        # # remove = [artist for artist in pickable_artists if     artist.contains(event)[0]]
        # remove = [artist for artist in X if artist.contains(event)[0]]
        #
        # if not remove:
        #     # add a pt
        #     x, y = ax.transData.inverted().transform_point([event.x,     event.y])
        #     pt, = ax.plot(x, y, 'o', picker=5)
        #     pickable_artists.append(pt)
        # else:
        #     for artist in remove:
        #         artist.remove()
        # plt.draw()
        # plt.draw_idle()

        xs, ys, zs = graph._offsets3d
        print(xs[ind[0]])
        print(ys[ind[0]])
        print(zs[ind[0]])
        print(dir(artist))

        # xs[ind[0]] = 0.5
        # ys[ind[0]] = 0.5
        # zs[ind[0]] = 0.5   
        # graph._offsets3d = (xs, ys, zs)

        # print(artist.get_facecolor())
        # artist.set_facecolor('red')
        graph._facecolors[ind, :] = (1, 0, 0, 1)

        plt.draw()

    def forceUpdate(event):
        global graph
        graph.changed()

    fig.canvas.mpl_connect('motion_notify_event', onMouseMotion)  # on mouse motion    
    fig.canvas.mpl_connect('button_press_event', onclick)
    fig.canvas.mpl_connect('pick_event', onpick)
    fig.canvas.mpl_connect('draw_event', forceUpdate)

    plt.tight_layout()

    plt.show()

Answer 1

OK, I have got at least a partial solution for you, without the rectangle selection, but you can select multiple points and delete them with one key_event.

To change the color, you need to change graph._facecolor3d, the hint was in this bug report about set_facecolor not setting _facecolor3d.

It might also be a good idea to rewrite your function as a class to get rid of any needed global variables.

My solution has parts that are not exactly pretty, I need to redraw the figure after removing data points, I couldn't get removing and updating to work. Also (see EDIT 2 below). I have not yet implemented what happens if the last data point is removed.

The reason why your function on_key(event) didn't work was easy: you forgot to connect it.

So this is a solution that should satisfy objectives (i) and (ii):

import matplotlib.pyplot as plt, numpy as np
from mpl_toolkits.mplot3d import proj3d

class Class3DDataVisualizer:    
    def __init__(self, X, ids, subindus, drawNew = True):

        self.X = X;
        self.ids = ids
        self.subindus = subindus

        self.disconnect = False
        self.ind = []
        self.label = None

        if drawNew:        
            self.fig = plt.figure(figsize = (7,5))
        else:
            self.fig.delaxes(self.ax)
        self.ax = self.fig.add_subplot(111, projection = '3d')
        self.graph  = self.ax.scatter(self.X[:, 0], self.X[:, 1], self.X[:, 2], depthshade = False, picker = True, facecolors=np.repeat([[0,0,1,1]],X.shape[0], axis=0) )         
        if drawNew and not self.disconnect:
            self.fig.canvas.mpl_connect('motion_notify_event', lambda event: self.onMouseMotion(event))  # on mouse motion    
            self.fig.canvas.mpl_connect('pick_event', lambda event: self.onpick(event))
            self.fig.canvas.mpl_connect('key_press_event', lambda event: self.on_key(event))

        self.fig.tight_layout()
        self.fig.show()


    def distance(self, point, event):
        """Return distance between mouse position and given data point

        Args:
            point (np.array): np.array of shape (3,), with x,y,z in data coords
            event (MouseEvent): mouse event (which contains mouse position in .x and .xdata)
        Returns:
            distance (np.float64): distance (in screen coords) between mouse pos and data point
        """
        assert point.shape == (3,), "distance: point.shape is wrong: %s, must be (3,)" % point.shape

        # Project 3d data space to 2d data space
        x2, y2, _ = proj3d.proj_transform(point[0], point[1], point[2], plt.gca().get_proj())
        # Convert 2d data space to 2d screen space
        x3, y3 = self.ax.transData.transform((x2, y2))

        return np.sqrt ((x3 - event.x)**2 + (y3 - event.y)**2)


    def calcClosestDatapoint(self, event):
        """"Calculate which data point is closest to the mouse position.

        Args:
            X (np.array) - array of points, of shape (numPoints, 3)
            event (MouseEvent) - mouse event (containing mouse position)
        Returns:
            smallestIndex (int) - the index (into the array of points X) of the element closest to the mouse position
        """
        distances = [self.distance (self.X[i, 0:3], event) for i in range(self.X.shape[0])]
        return np.argmin(distances)


    def annotatePlot(self, index):
        """Create popover label in 3d chart

        Args:
            X (np.array) - array of points, of shape (numPoints, 3)
            index (int) - index (into points array X) of item which should be printed
        Returns:
            None
        """
        # If we have previously displayed another label, remove it first
        if self.label is not None:
            self.label.remove()
        # Get data point from array of points X, at position index
        x2, y2, _ = proj3d.proj_transform(self.X[index, 0], self.X[index, 1], self.X[index, 2], self.ax.get_proj())
        self.label = plt.annotate( self.ids[index],
            xy = (x2, y2), xytext = (-20, 20), textcoords = 'offset points', ha = 'right', va = 'bottom',
            bbox = dict(boxstyle = 'round,pad=0.5', fc = 'yellow', alpha = 0.5),
            arrowprops = dict(arrowstyle = '->', connectionstyle = 'arc3,rad=0'))
        self.fig.canvas.draw()


    def onMouseMotion(self, event):
        """Event that is triggered when mouse is moved. Shows text annotation over data point closest to mouse."""
        closestIndex = self.calcClosestDatapoint(event)
        self.annotatePlot (closestIndex) 


    def on_key(self, event):
        """
        Function to be bound to the key press event
        If the key pressed is delete and there is a picked object,
        remove that object from the canvas
        """
        if event.key == u'delete':
            if self.ind:
                self.X = np.delete(self.X, self.ind, axis=0)
                self.ids = np.delete(ids, self.ind, axis=0)
                self.__init__(self.X, self.ids, self.subindus, False)
            else:
                print('nothing selected')

    def onpick(self, event):
        self.ind.append(event.ind)
        self.graph._facecolor3d[event.ind] = [1,0,0,1]



#Datapoints
Y = np.array([[ 4.82250000e+01,  1.20276889e-03,  9.14501289e-01], [ 6.17564688e+01,  5.95020883e-02, -1.56770827e+00], [ 4.55139000e+01,  9.13454423e-02, -8.12277299e+00], [3,  8, -8.12277299e+00]])
#Annotations
ids = ['a', 'b', 'c', 'd']

subindustries =  'example'

Class3DDataVisualizer(Y, ids, subindustries)

To implement your rectangular selection you would have to override what currently happens during dragging (rotating the 3D plot) or an easier solution would be to define your rectangle with two consecutive clicks.

Then use the proj3d.proj_transform to find which data is inside that rectangle, find the index of said data and recolor it using the self.graph._facecolor3d[idx] function and fill self.ind with those indices, after which hitting delete will take care of deleting all data which is specified by self.ind.

EDIT: I added two lines in the __init__ that remove the ax/subplot before adding a new one after data points are deleted. I noticed plot interactions were becoming slow after a few data points were removed as the figure was just plotting over each subplot.

EDIT 2: I found out how you can modify your data instead of redrawing the whole plot, as mention in this answer you will have to modify _offsets3d, which weirdly return a tuple for x and y, but an array for z.

You can modify it using

(x,y,z) = self.graph._offsets3d # or event.artist._offsets3d
xNew = x[:int(idx)] + x[int(idx)+1:]
yNew = y[:int(idx)] + y[int(idx)+1:]
z = np.delete(z, int(idx))
self.graph._offsets3d = (xNew,yNew,z) # or event.artist._offsets3d

But then you'll run into problem with deleting several data points in a loop because the indices that you stored before won't be applicable after the first loop, you'll have to update the _facecolor3d, the list of labels... so I chose to keep the answer as is, because just redrawing the plot with new data seems easier and cleaner than that.