Python/Matplotlib 带圆形热图添加尺寸图例与节点标注方法

以下代码完全适配你现有绘图逻辑，同时实现尺寸图例、节点双标注两个需求，同时修复原代码中未定义legend对象直接调用导致的运行报错：

from matplotlib.legend_handler import HandlerPatch
from matplotlib.collections import PatchCollection
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import pandas as pd
import numpy as np

df_hm_nodes = pd.DataFrame(
{
"abandonment ratio":[3.837299, 15.190000, 21.260000, 12.1600000, 14.710000],
"Total users": [1303, 1050, 105, 64, 34],
"Type of node":["Main node", "Standard node", "Standard node", "Standard node", "Type 45"],
"Name of node":["Node 1", "Node 3", "Node 4", "Node 64", "Node 45.1"]
}) 

cmap = "RdYlGn_r"

ylabels = df_hm_nodes["Name of node"].unique().tolist()
xlabels = df_hm_nodes["Type of node"].unique().tolist()
xn = len(xlabels)
yn = len(ylabels)    

s = df_hm_nodes["Total users"].values
c = df_hm_nodes["abandonment ratio"].values

# 生成节点对应x轴分类的坐标索引
x_mapping = {label:idx for idx, label in enumerate(xlabels)}
x_coords = [x_mapping[node_type] for node_type in df_hm_nodes["Type of node"]]

fig, ax = plt.subplots(figsize=(30,20))
ax.set_facecolor('#cecece')
ax.set_xlim(-0.5, xn-0.5)
ax.set_ylim(-0.5, yn-0.5)
ax.set(xticks=np.arange(xn), yticks=np.arange(yn), yticklabels=ylabels)
ax.set_xticklabels(xlabels, rotation='vertical')
ax.set_xticks(np.arange(xn)-0.5, minor=True)
ax.set_yticks(np.arange(yn)-0.5, minor=True)
ax.grid(which='minor')
ax.set_aspect("equal", "box")

R = s/s.max()/2
# 修正原代码所有圆形固定在x=0的问题，匹配对应分类列
circles = [plt.Circle((x_coords[i], i), radius=r) for i, r in enumerate(R)]
col = PatchCollection(circles, array=c, cmap=cmap)
sc=ax.add_collection(col)

# 添加节点双标注：圆形顶部居中显示放弃率、总用户数
for x, y, r, abandon_rate, total_user in zip(x_coords, range(yn), R, c, s):
    ax.text(
        x, y + r + 0.03, 
        f"{abandon_rate:.1f}%\n{int(total_user)}",
        ha='center', va='bottom', fontsize=12
    )

# 生成尺寸图例，匹配主图圆形缩放规则
legend_size_values = [50, 500, 1500]
legend_patches = []
for sz in legend_size_values:
    patch_r = sz / s.max() / 2
    legend_patches.append(mpatches.Circle((0,0), radius=patch_r, facecolor='#868686', edgecolor='white'))

size_legend = ax.legend(
    handles=legend_patches,
    labels=[f"{sz} 位用户" for sz in legend_size_values],
    loc='upper left',
    bbox_to_anchor=(1.03, 0.58),
    title='节点总用户数',
    handler_map={
        mpatches.Circle: HandlerPatch(lambda x,y,r,**kwargs: mpatches.Circle((x,y), r, **kwargs))
    }
)
ax.add_artist(size_legend)
plt.setp(size_legend.get_title(),fontsize='large')

# 颜色条设置，预留右侧图例空间
cbar = fig.colorbar(col, pad=0.18)
cbar.set_label('Abandonment rate', rotation=270, size=12, labelpad=20)

# 调整布局避免元素被截断
plt.subplots_adjust(right=0.8)
plt.show()

自定义调整说明 #

• 尺寸图例的档位可直接修改legend_size_values列表自定义，缩放逻辑和主图完全一致，不需要额外调整换算参数
• 如果仅需要显示放弃率标注，把标注部分的f字符串修改为f"{abandon_rate:.1f}%"即可
• 标注的偏移量、字体大小可根据实际输出效果微调，当前参数适配你设置的30*20英寸画布尺寸
• 尺寸图例和颜色条的位置可通过修改对应bbox_to_anchor、pad参数调整，避免遮挡主图内容
• 如果你的实际业务有自定义的节点x轴放置规则，直接修改x_coords的生成逻辑即可，不影响其他功能

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