How to add or update the color of a 3D curve efficiently (without for-loop)?

Question

Questions:

1. Is there a way to add color to a 3D curve plot more effectively than I am currently doing (e.g. avoiding the loop)?
2. Is there a way to update the color of a graph without plotting it all over again?
3. Is there a way to update one or more coordinate vectors without plotting it all over again?

Explanation:

I am making plots of 3D curves with color similar to the one in the example below. The number of points in my actual graphs is high, like n > 10000, which makes this approach very slow.

Also, the coordinates and colors are associated with the fixed number of points, which makes up the curves depend on another variable (i.e. time). Therefore I would like to plot the points once and then update one or all of the vectors (x, y, z, c) to speed up the plotting of each time step. EDIT: In Matlab you can do something like the following to update the coordinates of vertices/points already plotted. The example is with patches and not line segments, but the concept is the same:

v = [2 4; 2 8; 8 4; 5 0; 5 2; 8 0];
f = [1 2 3; 4 5 6];
col = [0; 1];
figure
p = patch('Faces',f,'Vertices',v,'FaceVertexCData',col,'FaceColor','flat');
colorbar
pause(1)
v = [1 3; 1 9;9 1; 5 0; 5 2; 8 0];
p.Vertices = v; % this updates the point coordinates which I would like to do for the line segments in my python plot.  
drawnow

Code used:

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

n = 100
x = np.arange(n)
y = np.arange(n)
z = np.arange(n)
colors = np.arange(n)/n

fig = plt.figure()
ax = fig.gca(projection='3d')
for i in range(1,len(x)-1):
    ax.plot(x[i-1:i+1], \
            y[i-1:i+1], \
            z[i-1:i+1], \
            c = plt.cm.jet(colors[i]))
plt.show()

EDIT: regarding my question 1: After gboffi's suggestion I looked into matplotlib.collections.LineCollection and can now plot the line segments much faster. I also used this answer Matplotlib Line3DCollection for time-varying colors.

from mpl_toolkits.mplot3d.art3d import Line3DCollection
import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure()
ax = plt.axes(projection='3d')
n = 100
x = np.arange(n)/n
y = np.arange(n)/n
z = np.arange(n)/n
colors = np.arange(n)/n
# "points" is used to restructere the coordinates array such that the first
# dimension is the point number and the third contains x, y and z of the point
points = np.array([x, y, z]).T.reshape(-1, 1, 3)
# "segs" is an array of lines between consecutive points in "points". So
# there is one line less than the number of points. The first dimension is
# the line number, second dimension is the start and end point, and the
# third dimension is the coordinates of the point.
segs = np.concatenate([points[:-1], points[1:]], axis=1)
cmap = plt.get_cmap('jet')
segment_color = cmap(colors)
ax.add_collection(Line3DCollection(segs, colors=segment_color))
plt.show()

Answer 1

EDIT: As suggested in gboffi's comment, use a LineCollection. They're even faster than the below scatter approach.

---

If you can live without the line segments, you can batch plot all the data points using scatter, which allows different colors for each data point, and which is quite fast.

Here's a small demonstration:

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

n = 100
x = np.arange(n)
y = np.arange(n)
z = np.arange(n)
colors = np.arange(n)/n

fig = plt.figure(figsize=(12, 5))

ax = fig.add_subplot(1, 2, 1, projection='3d')
for i in range(1,len(x)-1):
    ax.plot(x[i-1:i+1], y[i-1:i+1], z[i-1:i+1], c=plt.cm.jet(colors[i]))
ax.set_xlim(0, n), ax.set_ylim(0, n), ax.set_zlim(0, n)

ax = fig.add_subplot(1, 2, 2, projection='3d')
ax.scatter(x, y, z, s=1, c=colors, cmap='jet')
ax.set_xlim(0, n), ax.set_ylim(0, n), ax.set_zlim(0, n)

plt.show()

I did a small test with n = 100, and n = 10000, and also timed both runs:

On the left side, there's your code plotting line segments, on the right side, there's the scatter version. It's up to you to decide, whether you need the line segments to visualize the desired effect you want to express with your plots. At least for n = 10000, I personally can hardly see a difference between both versions.

But, obvious is the computational time difference! You now can even afford to re-plot ALL data, when changing any of x, y, z, c.

Hope that helps!

Answer 2

I start from your last edit, with the required imports

In [1]: import numpy as np 
   ...: import matplotlib.pyplot as plt 
   ...: from mpl_toolkits import mplot3d 
   ...: from mpl_toolkits.mplot3d.art3d import Line3DCollection     

Next, prepare the LineCollection — note that I have specified a colormap in the call and retained a handle to the data structure (different from your code)

In [2]: fig = plt.figure() 
   ...: ax = plt.axes(projection='3d') 
   ...: n = 100 
   ...: x = np.arange(n)/n 
   ...: y = np.arange(n)/n 
   ...: z = np.arange(n)/n 
   ...: colors = np.arange(n)/n 
   ...: points = np.array([x, y, z]).T.reshape(-1, 1, 3) 
   ...: segs = np.concatenate([points[:-1], points[1:]], axis=1) 
   ...: lc = Line3DCollection(segs, cmap='plasma', linewidth=4) 
   ...: ax.add_collection(lc)                                                             
Out[2]: <mpl_toolkits.mplot3d.art3d.Line3DCollection at 0x7f0528cdec10>

As you can see, we have not specified anything about the colors and our line came out in uniform blue.

Our instance of LineCollection, lc, has the property array, that is automatically mapped to colors using the colormap we specified (or 'viridis'), so we use the appropriate setter to notify lc that we want to change array and eventually we notify the ax that we want to redraw lc.

In [3]: lc.set_array(colors) ; ax.draw_artist(lc)                                         

Of course you can change the color-mapped array when you want, as long as you redraw lc.

Re changing the x y z data, you can change the segments using the appropriate setter, i.e., lc.set_segments but I'll leave that out of this answer.

Last, you can add a colorbar as in plt.colorbar(lc) `