为一系列线图添加色条

Question

我有一系列的折线图，显示了两个变量（x和y）在不同的z值下的变化。通常我会像这样添加折线图和图例：

import matplotlib.pyplot as plt

fig = plt.figure()
ax  = fig.add_subplot(111)
# suppose mydata is a list of tuples containing (xs, ys, z) 
# where xs and ys are lists of x's and y's and z is a number. 
legns = []
for(xs,ys,z) in mydata:
   pl = ax.plot(xs,ys,color = (z,0,0))
   legns.append("z = %f"%(z))
ax.legends(legns) 
plt.show()

但是我的图太多了，图例会遮住图形。我更希望有一个颜色条，显示与颜色对应的z值。我在图库里找不到这样的东西，而且我尝试处理颜色条的所有方法都失败了。显然，我必须先创建一组图形，然后才能尝试添加颜色条。

有没有简单的方法可以做到这一点？谢谢。

编辑（澄清）：

我想做类似这样的事情：

import matplotlib.pyplot as plt
import matplotlib.cm     as cm

fig = plt.figure()
ax  = fig.add_subplot(111)
mycmap = cm.hot
# suppose mydata is a list of tuples containing (xs, ys, z) 
# where xs and ys are lists of x's and y's and z is a number between 0 and 1
plots = []
for(xs,ys,z) in mydata:
   pl = ax.plot(xs,ys,color = mycmap(z))
   plots.append(pl)
fig.colorbar(plots)
plt.show()

但根据Matplotlib的参考资料，这样做不行，因为一组图形不是“可映射的”，这是什么意思我也不太清楚。

我使用LineCollection创建了一个替代的绘图函数：

def myplot(ax,xs,ys,zs, cmap):
    plot = lc([zip(x,y) for (x,y) in zip(xs,ys)], cmap = cmap)
    plot.set_array(array(zs))
    x0,x1 = amin(xs),amax(xs)
    y0,y1 = amin(ys),amax(ys)
    ax.add_collection(plot)
    ax.set_xlim(x0,x1)
    ax.set_ylim(y0,y1)
    return plot

xs和ys是x和y坐标的列表，而zs是不同条件的列表，用来给每条线上色。不过这感觉有点像是临时拼凑的解决方案……我原以为会有更好的方法。我喜欢plt.plot()函数的灵活性。

Answer 1

这里有一个稍微简化的例子，灵感来自于Boris和Hooked的最佳回答（感谢他们的好主意！）：

1. 离散色条

离散色条的实现稍微复杂一些，因为通过mpl.cm.get_cmap()生成的色图并不能直接用作colorbar()的参数。我们需要生成一个虚拟的可映射对象，下面的代码展示了如何做到这一点：

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl

n_lines = 5
x = np.linspace(0, 10, 100)
y = np.sin(x[:, None] + np.pi * np.linspace(0, 1, n_lines))
c = np.arange(1, n_lines + 1)

cmap = mpl.cm.get_cmap('jet', n_lines)

fig, ax = plt.subplots(dpi=100)
# Make dummie mappable
dummie_cax = ax.scatter(c, c, c=c, cmap=cmap)
# Clear axis
ax.cla()
for i, yi in enumerate(y.T):
    ax.plot(x, yi, c=cmap(i))
fig.colorbar(dummie_cax, ticks=c)
plt.show();

这段代码会生成一个带有离散色条的图表：

---

2. 连续色条

连续色条的实现相对简单，因为mpl.cm.ScalarMappable()可以让我们得到一个可以用于colorbar()的“图像”。

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl


n_lines = 5
x = np.linspace(0, 10, 100)
y = np.sin(x[:, None] + np.pi * np.linspace(0, 1, n_lines))
c = np.arange(1, n_lines + 1)

norm = mpl.colors.Normalize(vmin=c.min(), vmax=c.max())
cmap = mpl.cm.ScalarMappable(norm=norm, cmap=mpl.cm.jet)
cmap.set_array([])

fig, ax = plt.subplots(dpi=100)
for i, yi in enumerate(y.T):
    ax.plot(x, yi, c=cmap.to_rgba(i + 1))
fig.colorbar(cmap, ticks=c)
plt.show();

这段代码会生成一个带有连续色条的图表：

[附注] 在这个例子中，我个人不太明白为什么cmap.set_array([])是必要的（否则会出现错误信息）。如果有人了解背后的原理，请留言评论一下 :)

Answer 2

这里有一种方法可以在使用plt.plot()的同时实现这个功能。基本上，你可以先画一个临时的图，然后从这个图中获取颜色条。

import matplotlib as mpl
import matplotlib.pyplot as plt

min, max = (-40, 30)
step = 10

# Setting up a colormap that's a simple transtion
mymap = mpl.colors.LinearSegmentedColormap.from_list('mycolors',['blue','red'])

# Using contourf to provide my colorbar info, then clearing the figure
Z = [[0,0],[0,0]]
levels = range(min,max+step,step)
CS3 = plt.contourf(Z, levels, cmap=mymap)
plt.clf()

# Plotting what I actually want
X=[[1,2],[1,2],[1,2],[1,2]]
Y=[[1,2],[1,3],[1,4],[1,5]]
Z=[-40,-20,0,30]
for x,y,z in zip(X,Y,Z):
    # setting rgb color based on z normalized to my range
    r = (float(z)-min)/(max-min)
    g = 0
    b = 1-r
    plt.plot(x,y,color=(r,g,b))
plt.colorbar(CS3) # using the colorbar info I got from contourf
plt.show()

虽然这样有点浪费，但也挺方便的。如果你要画多个图，这样做就不算太浪费，因为你可以直接调用plt.colorbar()而不需要重新生成相关信息。

Answer 3

(我知道这个问题有点老了，但是...) 要显示颜色条，我们需要一个 matplotlib.cm.ScalarMappable。而 plt.plot 画的是线条，这些线条不能直接用来生成颜色条。因此，我们需要先创建一个可以映射数值的对象。

好的。那么 ScalarMappable 的构造函数需要一个 cmap 和一个 norm 实例。这里的 norm 是用来把数据缩放到 0 到 1 的范围，而 cmap 是你已经接触过的，它会把 0 到 1 之间的数字转换成颜色。所以在你的情况下：

import matplotlib.pyplot as plt
sm = plt.cm.ScalarMappable(cmap=my_cmap, norm=plt.normalize(min=0, max=1))
plt.colorbar(sm)

因为你的数据已经在 0 到 1 的范围内了，所以你可以简化 sm 的创建为：

sm = plt.cm.ScalarMappable(cmap=my_cmap)

编辑: 对于 matplotlib 版本 1.2 或更高，代码变为：

import matplotlib.pyplot as plt
sm = plt.cm.ScalarMappable(cmap=my_cmap, norm=plt.normalize(vmin=0, vmax=1))
# fake up the array of the scalar mappable. Urgh...
sm._A = []
plt.colorbar(sm)

编辑: 对于 matplotlib 版本 1.3 或更高，代码变为：

import matplotlib.pyplot as plt
sm = plt.cm.ScalarMappable(cmap=my_cmap, norm=plt.Normalize(vmin=0, vmax=1))
# fake up the array of the scalar mappable. Urgh...
sm._A = []
plt.colorbar(sm)

编辑: 对于 matplotlib 版本 3.1 或更高，简化为：

import matplotlib.pyplot as plt
sm = plt.cm.ScalarMappable(cmap=my_cmap, norm=plt.Normalize(vmin=0, vmax=1))
plt.colorbar(sm)