如何根据第三维数值为Matplotlib线段设置颜色并添加色条

现有Python代码可生成多段独立线段，需求是根据数组z的元素值为对应线段着色，以z的最小值和最大值作为coolwarm色图的两端，并在图右侧添加色条对应颜色与z值的映射关系。原始代码如下：

import numpy as np
import pylab as pl
from matplotlib import collections as mc
from matplotlib.colors import to_rgba

lines = [
    [(0, 1), (1, 1)],
    [(2, 3), (3, 3)],
    [(1, 2), (1, 3)]
]

colors = [
    (1,0,0,1), # red
    (0,1,0,1), # green
    (0,0,1,1), # blue
]

z = [1,2,3]
c = np.array([(1, 0, 0, 1), (0, 1, 0, 1), (0, 0, 1, 1)])

lc = mc.LineCollection(lines, colors=('r','g','b'), linewidths=2)
fig, ax = pl.subplots()
ax.add_collection(lc)
ax.autoscale()
ax.margins(0.1)
fig.savefig('line.segments.svg', bbox_inches='tight', pad_inches = 0.05)

要实现需求，需借助Matplotlib的Normalize和ScalarMappable类完成z值到颜色的映射，具体步骤如下：

• 归一化z值：用Normalize将z值缩放到[0,1]区间，对应色图的两端
• 映射颜色：通过cm.ScalarMappable将归一化后的z值映射到coolwarm色图的颜色
• 添加色条：将ScalarMappable对象传入colorbar方法，生成对应色条
• 更新LineCollection：将映射后的颜色传入LineCollection的colors参数

修改后的完整代码：

import numpy as np
import pylab as pl
from matplotlib import collections as mc
from matplotlib.colors import Normalize
import matplotlib.cm as cm

lines = [
    [(0, 1), (1, 1)],
    [(2, 3), (3, 3)],
    [(1, 2), (1, 3)]
]

z = [1, 2, 3]

# 归一化z值，设置色图范围为z的最小和最大值
norm = Normalize(vmin=np.min(z), vmax=np.max(z))
# 创建颜色映射器，使用coolwarm色图
cmap = cm.ScalarMappable(norm=norm, cmap='coolwarm')
# 将每个z值映射为对应的RGBA颜色
line_colors = cmap.to_rgba(z)

# 创建LineCollection，使用映射后的颜色
lc = mc.LineCollection(lines, colors=line_colors, linewidths=2)

fig, ax = pl.subplots()
ax.add_collection(lc)
ax.autoscale()
ax.margins(0.1)

# 添加色条
fig.colorbar(cmap, ax=ax)

fig.savefig('line_segments_with_colorbar.svg', bbox_inches='tight', pad_inches=0.05)
pl.show()

关键代码说明： #

• Normalize(vmin=np.min(z), vmax=np.max(z))：指定颜色映射的数值范围，z的最小值对应色图一端，最大值对应另一端
• cm.ScalarMappable(norm=norm, cmap='coolwarm')：建立数值到颜色的映射关系，采用蓝红渐变的coolwarm色图
• cmap.to_rgba(z)：将z数组中的每个值转换为对应的RGBA颜色，直接传给LineCollection的colors参数
• fig.colorbar(cmap, ax=ax)：在图表右侧添加色条，直观展示z值与颜色的对应关系

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