22

I've been googling quite some time with no success ... maybe my keywords are just lousy. Anyway, suppose I have three 1D numpy.ndarrays of the same length I'd like to plot them in 3D as a trajectory. Moreover, I'd like to be able to do either of the following things:

  1. Change the colour of the line as a function of z
  2. Change the colour of the line as a function of time (i.e. the index in the arrays)

This demo has an example of making such a curve:

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()

enter image description here

But how do I achieve 1 or 2? Solutions to only one or the other are welcome! Thanks in advance.

Ian Hincks
  • 3,608
  • 3
  • 23
  • 20

2 Answers2

25

As with normal 2d plots, you cannot have a gradient of color along an ordinary line. However, you can do it with 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()

I liked @Junuxx's hack so I applied it here:

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-Color by index 2-Color lines by index

Community
  • 1
  • 1
askewchan
  • 45,161
  • 17
  • 118
  • 134
  • 1
    I think the first statement isn't exactly correct, see my answer. But +1 for demonstration of jet colormap including normalization. – Junuxx Mar 25 '13 at 14:52
  • Well, it is technically true, hence you must plot `N` `line3D` objects, but your hack is pretty awesome :) – askewchan Mar 25 '13 at 14:58
  • 2
    It seems like `int(255/N*i)` is required for the color spec. For me a float results in the max value for the colormap. – gauteh Apr 29 '14 at 11:41
  • 1
    @gauteh `cm.jet` accepts a float from `0.0` to `1.0` or an int from `0` to `255`, so if you use a float greater than `1.0`, it will be interpreted as `1.0`. So if you want to pass a float, just leave off the `255`: `plt.cm.jet(i/N)`. For me, `N` and `i` are ints so I multiplied by `255`. – askewchan Apr 29 '14 at 17:02
  • When I try using the hack for a 2D version of this scatter plot, I receive the following error : `Invalid RGBA argument`. Does anyone know why? – Rylan Schaeffer Feb 02 '20 at 19:40
  • I'm trying to iteratively add line plots to the same figure. Maybe this is the problem? – Rylan Schaeffer Feb 02 '20 at 19:41
16

You can plot every line segment separately, as shown below. This just loops over 6 predefined colors, since @askewchan's answer already demonstrates well how to use a colormap.

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])

enter image description here

Junuxx
  • 14,011
  • 5
  • 41
  • 71