如何在seaborn catplot中按规则添加彩色椭圆标记差异？

解决思路 #

先预处理数据计算每个Key的A-B差值，匹配你设定的颜色规则，再在catplot的对应位置绘制高亮椭圆。

完整代码实现 #

# Stack Overflow Example
import numpy as np, pandas as pd, seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from random import choice
from string import ascii_lowercase, digits

chars = ascii_lowercase + digits

np.random.seed(8)
t = pd.DataFrame(
    {
    'Key': [''.join(choice(chars) for _ in range(2)) for _ in range(5)]*2,
    'Value': np.random.uniform(low=1, high=10, size=(10,)), 
    'Type': ['A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B']
    }
)

# 步骤1：预处理数据，计算A-B差值并匹配颜色规则
# 转成宽格式，每个Key对应A和B的Value
df_pivot = t.pivot(index='Key', columns='Type', values='Value').reset_index()
# 计算差值
df_pivot['A-B'] = df_pivot['A'] - df_pivot['B']
# 根据规则映射颜色
def get_highlight_color(diff):
    if diff > 4:
        return 'green'
    elif diff < -1:
        return 'red'
    elif 0 <= diff <= 2:
        return 'blue'
    else:
        return 'none'  # 不符合规则的不高亮

df_pivot['color'] = df_pivot['A-B'].apply(get_highlight_color)

# 步骤2：绘制catplot
g = sns.catplot(data=t, x='Value', y='Key', hue='Type', palette="dark")
g.set(title="Stack Overflow Help Me")
ax = g.axes[0,0]

# 获取每个Key对应的y轴坐标（seaborn的catplot中y轴是分类变量，坐标从0开始递增）
y_coords = {key: idx for idx, key in enumerate(df_pivot['Key'])}

# 遍历每个Key，绘制高亮椭圆
for _, row in df_pivot.iterrows():
    key = row['Key']
    color = row['color']
    if color == 'none':
        continue
    y = y_coords[key]
    # 椭圆的参数：(中心x, 中心y), 宽度, 高度, 角度（这里水平椭圆，角度0）
    # 宽度可以根据数据范围调整，这里取x轴范围的1/5作为椭圆宽度，高度设为0.6（适配分类y轴的间距）
    x_min, x_max = ax.get_xlim()
    ellipse_width = (x_max - x_min) * 0.2
    ellipse = Ellipse(( (x_min + x_max)/2, y ), width=ellipse_width, height=0.6, 
                     edgecolor=color, facecolor='none', linewidth=2)
    ax.add_patch(ellipse)

plt.show()

代码说明 #

1. 数据预处理：通过pivot将长格式数据转为宽格式，方便计算每个Key的A-B差值，再用自定义函数匹配对应的高亮颜色。
2. 椭圆绘制：获取catplot的轴对象后，遍历每个Key，找到其在y轴上的坐标，用Ellipse绘制边缘高亮的椭圆，椭圆的大小可以根据实际图表调整宽度和高度参数。
3. 适配规则：不符合三个规则的Key不会添加高亮椭圆，你可以根据需要修改get_highlight_color函数调整规则。

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