如何在带mask的heatmap边缘绘制轮廓线？

要在mask的边缘绘制突出线条，核心是定位每行mask的分界位置，然后用Matplotlib的绘图工具在对应坐标绘制线条。以下是修改后的完整代码：

import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np

vmin = 0
vmax = 80
cmap = "viridis"
size = 40

# 生成随机矩阵
matrix = np.random.randint(vmin, vmax, size=(size, size))
np.random.seed(7)

# 生成mask：每行前prefix_length个为False，其余为True
mask = np.zeros((size, size), dtype=bool)
split_cols = []
for _ in range(size):
    prefix_length = np.random.randint(size)
    mask[_ , prefix_length:] = True
    split_cols.append(prefix_length)
split_cols = np.array(split_cols)
# 修正全False行的分界值（避免argmax返回0）
split_cols[np.all(~mask, axis=1)] = size

# 绘制热图
ax = sns.heatmap(matrix, vmin=vmin, vmax=vmax, cmap=cmap, mask=mask)

# 遍历每行，绘制mask边缘的垂直线
for i in range(size):
    col_pos = split_cols[i]
    # 热图行索引i对应的y轴上下边界（Seaborn热图y轴从上到下递减）
    y_upper = size - i
    y_lower = size - i - 1
    # 绘制红色垂直线，线宽可根据需求调整
    ax.plot([col_pos, col_pos], [y_lower, y_upper], color='red', linewidth=1.5)

plt.savefig("temp.png")
plt.show()

关键步骤说明： #

• 定位mask分界点：记录每行中mask从False切换到True的列索引，同时处理全False的行（将分界点设为矩阵宽度）。
• 坐标适配：Seaborn热图的y轴是从上到下递减的，因此行索引i对应的y轴范围是size-i（顶部）到size-i-1（底部）。
• 绘制边缘线：针对每行的分界列位置，绘制垂直方向的线条，颜色和线宽可根据需求自定义（比如改成黑色、调整线宽为2）。

如果你的mask是其他不规则形状（非每行前缀式），可以通过计算mask的边缘检测来定位边界：利用np.diff找出mask中值变化的位置，再提取对应坐标绘制线条。

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