如何根据某个变量改变数据点颜色

Question

我有两个变量（x和y），它们会随着时间（t）变化。我想把x和t画成图，并根据y的值来给刻度上色。比如说，y值最高的时候，刻度的颜色是深绿色；y值最低的时候，刻度的颜色是深红色；而中间的值则会在绿色和红色之间渐变。

请问这可以用Python的matplotlib库实现吗？

Answer 1

如果你想画线而不是点，可以看看这个例子，这里做了一些修改，用黑色和红色来表示一个函数的好点和坏点：

def plot(xx, yy, good):
    """Plot data

    Good parts are plotted as black, bad parts as red.

    Parameters
    ----------
    xx, yy : 1D arrays
        Data to plot.
    good : `numpy.ndarray`, boolean
        Boolean array indicating if point is good.
    """
    import numpy as np
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots()
    from matplotlib.colors import from_levels_and_colors
    from matplotlib.collections import LineCollection
    cmap, norm = from_levels_and_colors([0.0, 0.5, 1.5], ['red', 'black'])
    points = np.array([xx, yy]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)
    lines = LineCollection(segments, cmap=cmap, norm=norm)
    lines.set_array(good.astype(int))
    ax.add_collection(lines)
    plt.show()

Answer 2

这就是 matplotlib.pyplot.scatter 的用途。

如果你没有指定颜色映射，scatter 就会使用默认的颜色映射。要指定使用哪个颜色映射，可以用 cmap 这个参数（比如 cmap="jet"）。

这里有个简单的例子：

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np

# Generate data...
t = np.linspace(0, 2 * np.pi, 20)
x = np.sin(t)
y = np.cos(t)

plt.scatter(t, x, c=y, ec='k')
plt.show()

你还可以指定自定义的颜色映射和标准化方式

cmap, norm = mcolors.from_levels_and_colors([0, 2, 5, 6], ['red', 'green', 'blue'])
plt.scatter(x, y, c=t, cmap=cmap, norm=norm)