Adding colors to a 3d quiver plot in matplotlib

后端 未结 4 1090
栀梦
栀梦 2020-12-11 03:40

I want to have colors corresponding to a colormap in my 3d quiver plot. The 2d version of the plot has an optional array that is used to map colors to the arrows. How can I

4条回答
  •  遥遥无期
    2020-12-11 04:03

    Building on @tacaswell and @sytrus answers, here is an example of coloring a 3d quiver plot

    from mpl_toolkits.mplot3d import axes3d
    import matplotlib.pyplot as plt
    import numpy as np
    
    # Make the grid
    x, y, z = np.meshgrid(np.arange(-0.8, 1, 0.2),
                          np.arange(-0.8, 1, 0.2),
                          np.arange(-0.8, 1, 0.8))
    
    # Make the direction data for the arrows
    u = np.sin(np.pi * x) * np.cos(np.pi * y) * np.cos(np.pi * z)
    v = -np.cos(np.pi * x) * np.sin(np.pi * y) * np.cos(np.pi * z)
    w = (np.sqrt(2.0 / 3.0) * np.cos(np.pi * x) * np.cos(np.pi * y) *
         np.sin(np.pi * z))
    
    # Color by azimuthal angle
    c = np.arctan2(v, u)
    # Flatten and normalize
    c = (c.ravel() - c.min()) / c.ptp()
    # Repeat for each body line and two head lines
    c = np.concatenate((c, np.repeat(c, 2)))
    # Colormap
    c = plt.cm.hsv(c)
    
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    ax.quiver(x, y, z, u, v, w, colors=c, length=0.1, normalize=True)
    plt.show()
    

提交回复
热议问题