Updating quiver 3D points in matplotlib - python-3.x

Question

I am trying to use matplotlib to graph a 3D tree. I want a user to be able to drag each node around so that they can view the tree in any manner they choose. My plan is to eventually subclass in order to add data to the nodes and create the tree.
I am having trouble updating the graph while they are dragging the nodes around. I found a way to update the position of the scatter points using:
self.plot.set_offsets([x,y])
self.plot.set_3d_properties([z], 'z')
But cannot seem to find a way of also updating the position and vector of the quivers.
I am also having difficulties with retrieving the mouse position which currently I always find:
ecoor == {'elevation': 30.0, 'azimuth': -60.0}
EDIT: Thanks to ImportanceOfBeingErnest instead of trying to update the plot I removed it and made a new one.
# PlotNode._plot changes
self.plot.remove()
self.plot = axes.scatter([x], [y], [z])
# PlotInterNode._plot changes (both vector updates)
arlen=length(xc - x, yc - y, zc - z)
edge.remove()
edge = axes.quiver(x, y, z, xc - x, yc - y, zc - z, length=arlen,
arrow_length_ratio=.5/arlen, pivot='tail')
self.edges.update({childname:(childnode, edge)})
Here is the original code:
import math
from matplotlib import pyplot
from mpl_toolkits.mplot3d import Axes3D
class PlotNode(object):
"""
Class for graph node.
"""
def __init__(self, name, parent):
"""
Initializes PlotNode.
"""
self.__name__ = name
self.parent = parent
self.coordinates = self._make_coordinates()
x, y, z = self.coordinates
self.plot = axes.scatter([x], [y], [z])
self.pressed = False
self.move = False
def name(self):
return self.__name__
COOR = (0, 0, 0)
def _make_coordinates(self):
"""
Finds coordinates from a file or, if not exist, calculate new coordinates.
"""
if PlotNode.COOR[1] == PlotNode.COOR[0]:
PlotNode.COOR = (PlotNode.COOR[0] + 1, 0, 0)
elif PlotNode.COOR[2] == PlotNode.COOR[1]:
PlotNode.COOR = (PlotNode.COOR[0], PlotNode.COOR[1] + 1, 0)
else:
PlotNode.COOR = (PlotNode.COOR[0], PlotNode.COOR[1], PlotNode.COOR[2] + 1)
return (PlotNode.COOR[0], PlotNode.COOR[1], PlotNode.COOR[2])
def _plot(self):
"""
Plots node onto graph.
"""
x, y, z = self.coordinates
#updates the plot coordinates
self.plot.set_offsets([x,y])
self.plot.set_3d_properties([z], 'z')
#updates the parent
if self.parent:
self.parent._plot(self.name())
self.plot.figure.canvas.draw()
def press(self, event):
"""
Mouse press event.
"""
if event.inaxes != self.plot.axes or not self.plot.contains(event)[0]:
return False
self.pressed = True
axes.disable_mouse_rotation() #Make sure node moves instead of plot rotation
return True
def release(self, event):
"""
Mouse release event.
"""
if event.inaxes != self.plot.axes or not self.pressed:
return False
self.pressed = False
if self.move:
self.move = False
x, y, z = self.coordinates
ecoor = to_dict(axes.format_coord(event.xdata, event.ydata))
xe, ye, ze = ecoor.get('x', x), ecoor.get('y', y), ecoor.get('z', z)
self.coordinates = (xe, ye, ze)
self._plot()
else:
self.open()
axes.mouse_init() #Make plot rotation avaliable again
return True
def motion(self, event):
"""
Mouse motion event.
"""
if event.inaxes != self.plot.axes or not self.plot.contains(event)[0]:
return False
if not self.pressed:
return False
self.move = True
x, y, z = self.coordinates
ecoor = to_dict(axes.format_coord(event.xdata, event.ydata))
xe, ye, ze = ecoor.get('x', x), ecoor.get('y', y), ecoor.get('z', z)
self.coordinates = (xe, ye, ze)
self._plot()
return True
def open(self):
print('openned!') #to be changed
class PlotInterNode(PlotNode):
"""
Class for graph folder node.
"""
def __init__(self, name, parent=None):
"""
Initializes PlotDir.
"""
self.edges = {}
PlotNode.__init__(self, name, parent)
def _plot(self, childname=None):
"""
Plots node onto graph.
"""
if childname:
x, y, z = self.coordinates
childnode, edge = self.edges.get(childname)
xc, yc, zc = childnode.coordinates
##update the vector
arlen=length(xc - x, yc - y, zc - z)
##update the arrow length
else:
x, y, z = self.coordinates
for childname in self.edges:
_, edge = self.edges.get(childname)
##update the position of each edge
super()._plot()
self.plot.figure.canvas.draw()
def traverse(self):
"""
Generator that traverses the tree rooted at this node.
"""
yield self
for child in self.edges:
try:
for node in self.edges.get(child)[0].traverse():
yield node
except AttributeError:
yield self.edges.get(child)[0]
def select_node(root, event):
"""
Single event function to handle all node movement.
"""
if event.name == 'button_press_event':
event_fn = lambda self: PlotNode.press(self, event)
elif event.name == 'button_release_event':
event_fn = lambda self: PlotNode.release(self, event)
elif event.name == 'motion_notify_event':
event_fn = lambda self: PlotNode.motion(self, event)
for node in root.traverse():
if event_fn(node):
return #if act on node then end
select_ids = []
def connect_select(root):
"""
Connects select_node to events.
"""
select_ids.append(figure.canvas.mpl_connect('button_press_event', lambda event: select_node(root, event)))
select_ids.append(figure.canvas.mpl_connect('button_release_event', lambda event: select_node(root, event)))
select_ids.append(figure.canvas.mpl_connect('motion_notify_event', lambda event: select_node(root, event)))
def to_dict(string):
"""
Converts a string to a dictionary.
"""
dictionary = {}
for st in string.split(','):
st = st.strip().split('=')
if st[0] not in dictionary:
try:
dictionary.update({st[0]:float(st[1])})
except ValueError:
st[1] = st[1].split(' ')[0]
dictionary.update({st[0]:float(st[1])})
return dictionary
def length(x, y, z):
"""
Returns hypotenuse.
"""
ret = math.sqrt(math.pow(x, 2) + math.pow(y, 2) + math.pow(z, 2))
if not ret:
ret = 1
return ret
figure = pyplot.figure()
axes = figure.add_subplot(111, projection='3d')
root = PlotInterNode('root')
def make_children(node, child_type, number):
x, y, z = node.coordinates
for i in range(number):
child = child_type(str(i), node)
xc, yc, zc = child.coordinates
arlen = length(xc - x, yc - y, zc - z)
edge = axes.quiver(x, y, z, xc - x, yc - y, zc - z, length=arlen, arrow_length_ratio=.5/arlen, pivot='tail')
node.edges.update({child.name():(child, edge)})
def node_depth(node, depth):
if not depth:
make_children(node, PlotNode, 3)
else:
make_children(node, PlotInterNode, 3)
for child in node.edges:
node_depth(node.edges.get(child)[0], depth-1)
node_depth(root, 3)
connect_select(root)
pyplot.show()
EDIT: All I needed to retrieve the mouse position in 3D was add
axes.button_pressed = None
before I called
axes.format_coord(event.xdata, event.ydata)

Answer 1

Thanks to ImportanceOfBeingErnest instead of trying to update the plot I removed it and made a new one.
# PlotNode._plot changes
self.plot.remove()
self.plot = axes.scatter([x], [y], [z])
# PlotInterNode._plot changes (both vector updates)
arlen=length(xc - x, yc - y, zc - z)
edge.remove()
edge = axes.quiver(x, y, z, xc - x, yc - y, zc - z, length=arlen,
arrow_length_ratio=.5/arlen, pivot='tail')
self.edges.update({childname:(childnode, edge)})
For the second part:
All I needed to retrieve the mouse position in 3D was add
axes.button_pressed = None
before I called
axes.format_coord(event.xdata, event.ydata)