Matplotlib中如何绘制带颜色映射的连续渐变3D曲线

现有4个长度均为n的数组x、y、z和T，需要使用matplotlib绘制3D曲线：(x, y, z)为曲线上各点的三维坐标，T为每个点对应的属性值（例如点的温度），要求将T值映射为曲线的显示颜色。

初始测试代码如下：

import numpy as np
from matplotlib import pyplot as plt

fig = plt.figure()
ax = fig.gca(projection='3d')
n = 100
cmap = plt.get_cmap("bwr")
theta = np.linspace(-4 * np.pi, 4 * np.pi, n)
z = np.linspace(-2, 2, n)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
T = (2*np.random.rand(n) - 1)  # 所有T值范围在[-1, 1]

目前检索到的两种实现方案都存在明显不足：

• 散点绘制方案：结合cmap调用scatter方法实现颜色映射，代码如下：

  ax = plt.gca()
  ax.scatter(x, y, z, cmap=cmap, c=T)
  

  该方法输出的是离散散点集合，不是连续曲线，不符合需求。
• 分段纯色绘制方案：将曲线拆分为n-1个线段，循环为每个线段单独设置纯色，代码如下：

  t = (T - np.min(T))/(np.max(T)-np.min(T))  # 归一化T值
  for i in range(n-1):
      plt.plot(x[i:i+2], y[i:i+2], z[i:i+2], c=cmap(t[i]))
  

  该方法每个线段只有单一纯色，无法实现平滑的颜色渐变效果，且最后一个点的T值没有被使用，存在数据遗漏。

使用Line3DCollection创建支持逐点颜色映射的3D连续曲线，可完美解决上述问题，完整实现代码如下：

import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection

# 生成示例数据
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
n = 100
cmap = plt.get_cmap("bwr")
theta = np.linspace(-4 * np.pi, 4 * np.pi, n)
z = np.linspace(-2, 2, n)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
T = (2*np.random.rand(n) - 1)  # 取值范围[-1, 1]

# 整理为Line3DCollection要求的线段格式：相邻两点组成一个线段单元
points = np.array([x, y, z]).T.reshape(-1, 1, 3)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

# 生成3D线段集合并绑定颜色映射规则
line_collection = Line3DCollection(segments, cmap=cmap, linewidth=2)
# 传入T值作为颜色映射的数据源，自动完成归一化与颜色匹配
line_collection.set_array(T)
ax.add_collection(line_collection)

# 手动设置坐标轴范围，避免线段集合无法自动适配坐标轴导致的显示不全问题
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min(), y.max())
ax.set_zlim(z.min(), z.max())

# 添加颜色条，展示T值与颜色的对应关系
fig.colorbar(line_collection, ax=ax, label='属性值T')

plt.show()

该方案的优势：

• 输出为连续曲线，无离散散点的断裂感
• 颜色沿曲线平滑过渡，无分段纯色的生硬边界
• 所有采样点的T值全部参与颜色计算，无数据遗漏
• 原生支持颜色条组件，可直观展示属性值和颜色的映射关系

内容的提问来源于stack exchange，提问作者Carlos Adir