在heatmap上绘制带标签点：是否有更优实现方法？

问题描述 #

我需要在已有的热力图上绘制特定的带标签点，目前已通过一段代码实现预期效果，但想知道是否存在更简洁高效的实现方式。

现有实现代码：

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# 创建热力图数据
df = pd.DataFrame({
    'x': [1.1, 1.1, 2.1, 1.1, 2.1, 2.1],
    'y': [1.4, 2.4, 1.4, 3.4, 2.4, 3.4],
    'value': [9, 8, 1, 0, 6, 2]
})

# 转换为热力图所需的透视表格式
df = df.pivot(index='y', columns='x', values='value')

# 绘制热力图
fig, ax = plt.subplots()
sns.heatmap(df, annot=False, cmap='viridis')

# 待绘制的点数据
points = pd.DataFrame({
    'x': [1.8, 2.1],
    'y': [2.3, 1.0]
})

# 获取坐标轴标签和位置，手动转换坐标
x_labels = ax.get_xticklabels()
x_tick_start = x_labels[0].get_position()[0]
x_label_start = float(x_labels[0].get_text())
x_tick_end = x_labels[-1].get_position()[0]
x_label_end = float(x_labels[-1].get_text())
y_labels = ax.get_yticklabels()
y_tick_start = y_labels[0].get_position()[1]
y_label_start = float(y_labels[0].get_text())
y_tick_end = y_labels[-1].get_position()[1]
y_label_end = float(y_labels[-1].get_text())

# 转换点坐标到热力图网格坐标系
points['x_new'] = points['x'].apply(lambda x: ((x - x_label_start) * (x_tick_end - x_tick_start)) / (x_label_end - x_label_start) + x_tick_start)
points['y_new'] = points['y'].apply(lambda y: ((y - y_label_start) * (y_tick_end - y_tick_start)) / (y_label_end - y_label_start) + y_tick_start)

# 绘制散点
ax.scatter(x=points['x_new'], y=points['y_new'], color='black')

# 添加标签
for i, row in points.iterrows():
    plt.text(row['x_new']+0.04, row['y_new']+0.05, i+1)

plt.show()

更优解决方案 #

可以利用Matplotlib内置的坐标转换功能，或借助Numpy的线性插值实现坐标映射，避免手动计算的冗余步骤，代码更简洁可靠：

方法一：使用ax.transLimits快速转换坐标

ax.transLimits可直接将数据坐标转换为轴的归一化坐标，再映射到显示坐标，步骤更简洁：

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# 创建热力图数据
df = pd.DataFrame({
    'x': [1.1, 1.1, 2.1, 1.1, 2.1, 2.1],
    'y': [1.4, 2.4, 1.4, 3.4, 2.4, 3.4],
    'value': [9, 8, 1, 0, 6, 2]
})
df = df.pivot(index='y', columns='x', values='value')

# 绘制热力图
fig, ax = plt.subplots()
sns.heatmap(df, annot=False, cmap='viridis')

# 点数据（直接包含标签）
points = pd.DataFrame({'x': [1.8, 2.1], 'y': [2.3, 1.0], 'label': [1, 2]})

# 转换坐标：利用transLimits将数据坐标转为轴坐标
points['x_ax'] = points['x'].apply(lambda x: ax.transLimits.transform((x, 0))[0])
points['y_ax'] = points['y'].apply(lambda y: ax.transLimits.transform((0, y))[1])

# 绘制点和标签
for _, row in points.iterrows():
    ax.scatter(row['x_ax'], row['y_ax'], color='black')
    ax.text(row['x_ax'] + 0.02, row['y_ax'] + 0.02, row['label'], color='red', fontweight='bold')

plt.show()

方法二：基于Numpy线性插值映射坐标

通过热力图的x/y轴刻度值和对应位置，用线性插值快速完成坐标转换：

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

# 创建热力图数据
df = pd.DataFrame({
    'x': [1.1, 1.1, 2.1, 1.1, 2.1, 2.1],
    'y': [1.4, 2.4, 1.4, 3.4, 2.4, 3.4],
    'value': [9, 8, 1, 0, 6, 2]
})
df = df.pivot(index='y', columns='x', values='value')

# 绘制热力图
fig, ax = plt.subplots()
sns.heatmap(df, annot=False, cmap='viridis')

# 获取热力图的x/y轴刻度值和对应轴位置
x_vals = df.columns.values
x_pos = ax.get_xticks()
y_vals = df.index.values
y_pos = ax.get_yticks()

# 点数据
points = pd.DataFrame({'x': [1.8, 2.1], 'y': [2.3, 1.0], 'label': [1, 2]})

# 用Numpy线性插值转换坐标
points['x_ax'] = np.interp(points['x'], x_vals, x_pos)
points['y_ax'] = np.interp(points['y'], y_vals, y_pos)

# 绘制点和标签
for _, row in points.iterrows():
    ax.scatter(row['x_ax'], row['y_ax'], color='black', s=50)
    ax.text(row['x_ax'] + 0.05, row['y_ax'] + 0.05, row['label'], fontsize=12)

plt.show()

方案优势 #

• 代码更简洁，避免手动计算坐标转换的冗余逻辑
• 依赖内置函数实现转换，可靠性更高，不易出错
• 可读性更强，后续维护成本更低

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