如何在Python生成的Heatmap中添加矩形补丁高亮指定数值区域？

如下代码可生成热图输出，我希望添加矩形补丁来高亮数值处于0.4_{0.99以及-0.4}-0.99的区域：

plt.figure(figsize=(15,10))
mask = np.triu(np.ones_like(corr, dtype=np.bool)) 
sns.heatmap(corr,annot=True,fmt=".2f", mask=mask,cmap="YlGnBu");

解决方法 #

你可以遍历热图的所有单元格，判断数值是否落在目标区间，对符合要求的单元格添加矩形边框补丁即可，完整可运行代码如下：

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

plt.figure(figsize=(15,10))
mask = np.triu(np.ones_like(corr, dtype=np.bool)) 
# 存储热图坐标轴对象
ax = sns.heatmap(corr, annot=True, fmt=".2f", mask=mask,cm="YlGnBu")

# 遍历所有单元格
for i in range(corr.shape[0]):
    for j in range(corr.shape[1]):
        # 若corr为numpy数组，替换为corr[i,j]
        val = corr.iloc[i,j]
        # 判断数值是否在目标区间
        if (0.4 <= val <= 0.99) or (-0.99 <= val <= -0.4):
            # 添加矩形高亮边框，可自行调整颜色和线宽
            ax.add_patch(Rectangle((j,i), 1,1, edgecolor='red', lw=2, fill=False))

plt.show()

• 可修改edgecolor参数调整高亮边框的颜色，修改lw参数调整边框粗细
• 代码会自动跳过上三角被mask的区域，不需要额外做过滤处理

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