Plotting multiple 3D lines: data transformation

Question

I would like to plot multiple lines in a 3D plot in Python. My input consists of two n x 3 arrays, say pos1 and pos2, corresponding to two lists of 3D points (row = x,y,z coordinates). I need to plot a line connecting the ith point in pos1 to the ith point in pos2, for each i.

I have working code, but I am certain that it is terrible and that there is a much better way to implement this.

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np

# get plot set up
fig = plt.figure()
ax = fig.add_subplot(111, projection = '3d')

# minimal working example: two pairs.
pos1 = np.array([[0,0,0],[2,2,2]])
pos2 = np.array([[1,1,1],[3,3,3]])

xvals = np.transpose(np.array([pos1[:,0], pos2[:,0]]))
yvals = np.transpose(np.array([pos1[:,1], pos2[:,1]]))
zvals = np.transpose(([pos1[:,2], pos2[:,2]]))

N = len(pos1)

# call plot on each line
for i in range(N):
    ax.plot(xvals[i],yvals[i],zvals[i])

plt.show()

Specifically, I do not think it is necessary to define the three intermediate arrays xvals, yvals, zvals. Is there a cleaner alternative to this? Should I be passing a list to plot, for example?

Answer 1

I've never done 3D plotting before, so I can't say whether this would be the best method using the tools matplotlib has to offer. But you could make good use of the zip function.

pos1 = np.array([[0,0,0],[2,2,2]])
pos2 = np.array([[1,1,1],[3,3,3]])

for point_pairs in zip(pos1, pos2):
    xs, ys, zs = zip(*point_pairs)
    ax.plot(xs, ys, zs)