Цвет линии 3D параметрической кривой в python matplotlib.pyplot

Я уже давно занимаюсь поиском в Интернете без успеха... возможно, мои ключевые слова просто паршивые. В любом случае, предположим, что у меня есть три 1D numpy.ndarray той же длины, которые я хотел бы построить в 3D в качестве траектории. Более того, я хотел бы иметь возможность выполнить одно из следующих действий:

Измените цвет линии как функцию z
Измените цвет линии как функцию времени (т.е. индекс в массивах)

В этой демонстрации приведен пример такой кривой:

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

fig = plt.figure()
ax = fig.gca(projection='3d')
theta = np.linspace(-4 * np.pi, 4 * np.pi, 100)
z = np.linspace(-2, 2, 100)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
ax.plot(x, y, z)

plt.show()

Но как мне достичь 1 или 2? Разрешения только для одного или другого приветствуются! Спасибо заранее.

Ответ 1

Как и при нормальных 2d-графиках, вы не можете иметь градиент цвета вдоль обычной линии. Однако вы можете сделать это с помощью scatter:

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

fig = plt.figure()
ax = fig.gca(projection='3d')
theta = np.linspace(-4 * np.pi, 4 * np.pi, 100)
z = np.linspace(-2, 2, 100)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)

#1 colored by value of `z`
ax.scatter(x, y, z, c = plt.cm.jet(z/max(z))) 

#2 colored by index (same in this example since z is a linspace too)
N = len(z)
ax.scatter(x, y, z, c = plt.cm.jet(np.linspace(0,1,N)))

plt.show()

Мне понравился @Junuxx hack, поэтому я применил его здесь:

for i in xrange(N-1):
    ax.plot(x[i:i+2], y[i:i+2], z[i:i+2], color=plt.cm.jet(255*i/N))

Ответ 2

Вы можете построить каждый сегмент отдельно, как показано ниже. Это просто зацикливается на 6 предопределенных цветов, так как @askewchan answer уже хорошо показывает, как использовать цветную карту.

cols = 'rgbcmy'

for i in range(len(x)-1):
    ax.plot(x[i:i+2], y[i:i+2], z[i:i+2], color=cols[i%6])