CODEWITHSUNDEEP
python3dmatplotlib
thosphor
thosphor

Reputation: 2781

Plot stack of unit vectors in cylindrical coordinates - matplotlib

I have a python program that calculates angles for me and outputs them in a list.

What I would like to do is plot a stack of arrows that are unit vectors pointing in the direction of the angle. So I thought cylindrical coordinates would be best since they only have one angular coordinate.

I've tried pyplot.quiver but I don't think that can do anything in 3D, and a 3D line plot didn't work either.

Is there a way of doing this without laboriously converting each (length, height, angle) into a pair of vectors (a, b, c),(length*cos(angle), length*sin(angle), height)?

Upvotes: 0

Views: 1589

Answers (1)

wflynny
wflynny

Reputation: 18551

If you have a list of angles, you can easily calculate vectors associated with those angles using numpy.

import numpy as np
import matplotlib.pyplot as plt
angles = np.random.rand(100)

length = 1.
vectors_2d = np.vstack((length * np.cos(angles), length * np.sin(angles))).T

for x, y in vectors_2d:
    plt.plot([0, x], [0, y])
plt.show()

enter image description here

If you really want it in cylindrical instead of polar coords, then 

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
angles = np.random.rand(100)

length = 1.
heights = np.arange(len(angles))
vectors_3d = np.vstack((length * np.cos(angles), 
                        length * np.sin(angles), 
                        heights)).T

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
for x, y, z in vectors_3d:
    ax.plot([0, x], [0, y], zs=[z, z])
plt.show()

enter image description here

Edit: I know how to put arrows on plots using pyplot.quiver. However, I don't think mplot3d plays nicely with quiver. Maybe someone like @tcaswell can help out with a work around. But in 2D, you can do

import numpy as np
import matplotlib.pyplot as plt

angles = np.random.rand(100)
# Define coords for arrow tails (the origin)
x0, y0 = np.zeros(100), np.zeros(100)
# Define coords for arrow tips (cos/sin)
x, y = np.cos(angles), np.sin(angles)

# in case you want colored arrows
colors = 'bgrcmyk'
colors *= colors * (len(x0) / len(colors) + 1)
plt.quiver(x0, y0, x, y, color=colors[:len(x0)], scale=1) #scale sets the length
plt.show()

Upvotes: 1

Related Questions