Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

"set_UVC" equivilent for a 3D quiver plot in matplotlib

To update a regular 2D quiver plot you can use .set_UVC() to directly set the x and y data.

The 3D equivalent to the quiver is the Axes3D.quiver(), this however does not seem to have an equivalent to .set_UVC(). How can the data updated? The segments does seem to contain the data from the quiver arrows however I don't see how the input data and the segments correlates.

I can delete the quiver plot then re-plot a new one however that is inefficient and would impact the performance, I want know if there is a way to set the data directly.

like image 543
AceLewis Avatar asked Feb 21 '18 17:02

AceLewis


2 Answers

If you have a look at the Line3DCollection code, you'll see that not too many methods from the LineCollection class have been overloaded. For you the important one is set_segments(), which looks like this:

def set_segments(self, segments):
    '''                                                                                                                
    Set 3D segments                                                                                                    
    '''
    self._segments3d = np.asanyarray(segments)
    LineCollection.set_segments(self, [])

So when set_segments() is called, the segments are actually stored in self._segments3d and the LineCollection's set_segments() method is called with an empty list. Line3DCollection then takes care of its own segments list in an overloaded draw() method. Two things to notice:

  1. Even if you follow the matplotlib example and use numpy.meshgrids to define your quiver coordinates, within _segments3d the coordinates are stored in an array of shape (N,2,3), where N is the amount of points and its content is basically [[[x0,y0,z0],[u0,v0,w0]],[[x1,y1,z1],[u1,v1,w1]],...], so you probably have to manipulate your data to fit that format.
  2. Apparently you cannot assign the new values directly to _segments3d. At least for me that resulted in the figure not updating properly -- you have to go through set_segments(). However, you can access the previous values by reading out _segments3d, if you, for instance, just want to alter some of the coordinates.

Here still some arbitrary example which I produced to test all the things I just explained:

from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from matplotlib.animation import FuncAnimation
import numpy as np

fig = plt.figure()
ax = fig.gca(projection='3d')

num_frames = 50
theta = np.linspace(0,2*np.pi, 10, endpoint=False)
r = np.arange(1,2.1)
z = np.arange(-2,2.1,1)


def compute_segs(i):
    offset = 2*i*np.pi/num_frames
    theta2,r2, z2 = np.meshgrid(theta+offset,r,z)

    x = r2*np.cos(theta2)
    y = r2*np.sin(theta2)

    u = x+0.2*np.cos(4*theta2)
    v = y
    w = z2+0.2*np.sign(z2)*np.sin(4*theta2)

    return x,y,z2,u,v,w


segs = compute_segs(0)
cols = ['b' for x in segs[0].ravel()]
cols[0] = 'r'
quivers = ax.quiver(*segs, length=0.1, colors = cols, normalize=True)


ax.set_xlim([-3,3])
ax.set_ylim([-3,3])
ax.set_zlim([-3,3])
def animate(i):

    segs = np.array(compute_segs(i)).reshape(6,-1)

    new_segs = [[[x,y,z],[u,v,w]] for x,y,z,u,v,w in zip(*segs.tolist())]
    quivers.set_segments(new_segs)
    return quivers


ani = FuncAnimation(fig, animate, frames = num_frames, interval = 30, blit=False)
ani.save('update_3d_quiver.gif', writer='imagemagick')

plt.show()

...and the result looks like this:

result of the above code

Hope this helps.

like image 133
Thomas Kühn Avatar answered Nov 03 '22 01:11

Thomas Kühn


To elaborate on Thomas Kühn's answer that I have accepted, if you have your grid points and the data for the quivers that you want to convert to segments you can use the following function.

def quiver_data_to_segments(X, Y, Z, u, v, w, length=1):
    segments = (X, Y, Z, X+v*length, Y+u*length, Z+w*length)
    segments = np.array(segments).reshape(6,-1)
    return [[[x, y, z], [u, v, w]] for x, y, z, u, v, w in zip(*list(segments))]

You can then update the segments of your plot using the output, this way also allows you to use the length to specify the normalisation length of the arrows, it is usually best to set this to a multiple of the mean or maximum of the input data (np.sqrt(vv**2 + uu**2 + ww**2)).

segments = quiver_data_to_segments(X, Y, Z, uu, vv, ww, length=5)
quiver_plot.set_segments(segments)
like image 22
AceLewis Avatar answered Nov 03 '22 00:11

AceLewis