在 Matplotlib 中创建多色线条

简介

在本实验中，我们将学习如何在 Matplotlib 中创建多色线条。我们将使用 LineCollection 函数创建一组线段，并根据它们的导数分别为其着色。我们还将学习如何使用边界规范为线段着色。

虚拟机使用提示

虚拟机启动完成后，点击左上角切换到 笔记本 标签页，以访问 Jupyter Notebook 进行练习。

有时，你可能需要等待几秒钟让 Jupyter Notebook 完成加载。由于 Jupyter Notebook 的限制，操作验证无法自动化。

如果你在学习过程中遇到问题，随时向 Labby 提问。课程结束后提供反馈，我们将立即为你解决问题。

导入库

我们将首先为本实验导入必要的库。我们将导入 matplotlib.pyplot 和 numpy。

import matplotlib.pyplot as plt
import numpy as np

创建数据

我们将创建一个 NumPy 数组 x，它包含在 0 到 3π之间均匀分布的 500 个值。我们还将创建另一个 NumPy 数组 y，它包含 x 中值的正弦值。最后，我们将创建一个 NumPy 数组 dydx，它包含 y 的一阶导数。

x = np.linspace(0, 3 * np.pi, 500)
y = np.sin(x)
dydx = np.cos(0.5 * (x[:-1] + x[1:]))

创建线段

我们将创建一组线段，以便能够分别为它们上色。我们将使用 NumPy 的 concatenate 函数，沿着第二个轴连接两个数组 points[:-1] 和 points[1:]。然后，我们会将结果数组重塑为一个 N x 1 x 2 的数组，这样我们就可以轻松地将点堆叠在一起以获得线段。用于线集合的线段数组需要是 (线段数量) x (每条线段的点数) x 2（用于 x 和 y）。

points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

创建连续规范

我们将创建一个连续规范，用于将数据点映射到颜色。我们将使用 matplotlib.pyplot 中的 Normalize 函数，将 dydx 的值在其最小值和最大值之间进行归一化。然后，我们将使用 LineCollection 函数创建一组线段，并根据它们的导数分别为其着色。我们将使用 set_array 函数来设置用于颜色映射的值。

norm = plt.Normalize(dydx.min(), dydx.max())
lc = LineCollection(segments, cmap='viridis', norm=norm)
lc.set_array(dydx)
lc.set_linewidth(2)

创建一个颜色条

我们将创建一个颜色条来展示颜色与 dydx 值之间的映射关系。我们会使用 matplotlib.pyplot 中的 colorbar 函数来创建一个颜色条。

line = plt.gca().add_collection(lc)
plt.colorbar(line)

使用边界规范

我们将改为使用边界规范来为线段着色。我们将创建一个包含三种颜色（红色、绿色和蓝色）的 ListedColormap。然后，我们将创建一个边界为 -1、-0.5、0.5 和 1 的 BoundaryNorm，以及上述 ListedColormap。我们将使用 set_array 函数来设置用于颜色映射的值。

cmap = ListedColormap(['r', 'g', 'b'])
norm = BoundaryNorm([-1, -0.5, 0.5, 1], cmap.N)
lc = LineCollection(segments, cmap=cmap, norm=norm)
lc.set_array(dydx)
lc.set_linewidth(2)

创建一个子图

我们将创建一个子图来展示带颜色的线段。我们会使用 matplotlib.pyplot 中的 subplots 函数创建一个 2x1 的子图网格，并使用 sharex 和 sharey 参数在子图之间共享 x 轴和 y 轴。

fig, axs = plt.subplots(2, 1, sharex=True, sharey=True)
line = axs[0].add_collection(lc)
fig.colorbar(line, ax=axs[0])

设置坐标轴范围

我们将为子图设置 x 轴和 y 轴的范围。

axs[0].set_xlim(x.min(), x.max())
axs[0].set_ylim(-1.1, 1.1)

显示绘图

我们将使用 matplotlib.pyplot 中的 show 函数来显示绘图。

plt.show()

总结

在本实验中，我们学习了如何在 Matplotlib 中创建多色线条。我们使用 LineCollection 函数创建一组线段，并根据它们的导数分别为其着色。我们还学习了如何使用边界规范为线段着色。我们使用 Normalize 函数将 dydx 值在其最小值和最大值之间进行归一化，并使用 ListedColormap 函数创建一个包含三种颜色（红色、绿色和蓝色）的颜色映射。我们使用 BoundaryNorm 函数创建一个边界为 -1、-0.5、0.5 和 1 的边界规范，以及上述 ListedColormap。最后，我们使用 subplots 函数创建一个 2x1 的子图网格，并使用 sharex 和 sharey 参数在子图之间共享 x 轴和 y 轴。然后，我们使用 show 函数显示绘图。