使用Python Matplotlib根据数组为绘制线条设置对应颜色编码

原代码的问题是在嵌套循环中重复绘制了通长的水平/垂直线，既无法实现分段着色，也没有建立数组P和线条的对应映射关系，修改逻辑如下：

• 移除嵌套循环中重复绘制通长线条的代码，仅保留黑色方块的绘制逻辑
• 基于P的取值范围做颜色归一化，使用和预期效果匹配的viridis颜色映射（低数值对应紫色、高数值对应黄色）
• 按顺序逐段绘制水平、垂直线段，每段对应P中一个数值取色：
  • 先绘制3根水平线，每根线拆为2个可见分段（避开黑色方块覆盖区域），按y从低到高、同一条线x从左到右的顺序，对应P前6个值
  • 再绘制3根垂直线，每根线拆为2个可见分段，按x从左到右、同一条线y从低到高的顺序，对应P后6个值
• 适当调大线条宽度，匹配预期显示效果

修改后完整代码 #

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import numpy as np

fig,ax = plt.subplots(1)
n=3
# 仅绘制黑色方块，不在嵌套循环中绘制通长线条
for i in range(0,n):
    for j in range(0,n):
        rect = mpl.patches.Rectangle((200+200*i,200+200*j),10*n, 10*n, linewidth=1, edgecolor='black', facecolor='black')
        ax.add_patch(rect)

# 定义颜色数组与映射规则
P=np.array([[1.9],
       [4.9],
       [6.1],
       [8.2],
       [1.8],
       [5.8],
       [9.7],
       [7.3],
       [8.9],
       [2.5],
       [9.9],
       [0.7]])
norm = mpl.colors.Normalize(vmin=P.min(), vmax=P.max())
cmap = plt.get_cmap('viridis')
idx = 0

# 逐段绘制水平线，对应P前6个值
for i in range(n):
    y_pos = 200 + 200*i + 5*n
    # 跳过黑块覆盖区域，绘制两个独立分段
    for x_start in [200 + 10*n, 400 + 10*n]:
        x_end = x_start + 200 - 10*n
        line_color = cmap(norm(P[idx, 0]))
        ax.hlines(y_pos, x_start, x_end, color=line_color, zorder=0, linewidth=3)
        idx += 1

# 逐段绘制垂直线，对应P后6个值
for j in range(n):
    x_pos = 200 + 200*j +5*n
    # 跳过黑块覆盖区域，绘制两个独立分段
    for y_start in [200 + 10*n, 400 + 10*n]:
        y_end = y_start + 200 - 10*n
        line_color = cmap(norm(P[idx, 0]))
        ax.vlines(x_pos, y_start, y_end, color=line_color, zorder=0, linewidth=3)
        idx +=1

ax.set_xlim(left = 0, right = 220*n)
ax.set_ylim(bottom = 0, top = 220*n)
ax.set_aspect('equal')
plt.show()

运行代码后即可得到预期效果：所有线条分段严格按照P的顺序映射颜色，黑色方块样式保持不变，颜色从低数值的深紫色平滑过渡到高数值的亮黄色。

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