循环生成4×2并排Seaborn热力图问题求助（含代码修正）

我尝试遍历small、medium、large、extra large四种公司规模，为percent_engaged和avg_sales两个指标创建4×2布局的热力图——每个公司规模对应两个并排的指标热力图，同时要求公司规模按从小到大排序，但当前热力图完全无显示，附上代码与样本数据寻求解决。

原代码 #

fig, ax= plt.subplots(4,2,figsize(30,50))
metrics = ['avg_sales', 'percent_engaged']
for company_size,  ax in zip(df.company_size.unique(),axes.flat):
   df2 = df[df.company_size==company_size]
   for metric in metrics:
   sns.heatmap(df2.pivot(index='months_since_purchase', columns='year', values='metrics'),annot=True,cmap='YlBnBu',linewidth=2.5, linecolor='white', cbar=False, ax=ax)
ax.set_title("Purchases per company for: " + company_size)
ax.set_xlabel("Year")
ax.set_ylabel("Months Since First Purchase")
ax.tick_params(axis='both', which='major', labelsize=12)
plt.show()

样本数据说明 #

样本数据列：['year', 'company_size', 'months_since_purchase', 'percent_engaged', 'avg_sales']

样本数据示例：

2019-01-01  small   0   1.00    2.00
2019-01-01  small   1   0.90    3.00
2019-01-01  small   2   0.86    2.94
...（其余数据省略）

问题排查与修复方案 #

核心错误点 #

1. 变量名冲突与引用错误：plt.subplots返回的子图对象命名为ax，但循环中又用ax作为迭代变量，同时错误引用了未定义的axes
2. 指标参数错误：pivot中values='metrics'是字符串，实际应引用变量metric
3. 循环逻辑不匹配：未对应4×2的布局结构，无法正确分配子图位置
4. 公司规模无序：df.company_size.unique()返回结果不保证顺序，需手动指定排序
5. 语法错误：for metric in metrics:后缺少缩进，代码块不合法
6. figsize调用错误：figsize(30,50)需改为figsize=(30,50)

修正后完整代码 #

import matplotlib.pyplot as plt
import seaborn as sns

# 指定公司规模的排序顺序
size_order = ['small', 'medium', 'large', 'extra large']
metrics = ['avg_sales', 'percent_engaged']

# 创建4行2列的子图布局，修正figsize调用方式
fig, axes = plt.subplots(4, 2, figsize=(30, 50))

# 遍历公司规模（行）和指标（列），对应子图位置
for row, company_size in enumerate(size_order):
    df_filtered = df[df['company_size'] == company_size]
    for col, metric in enumerate(metrics):
        # 生成对应指标的透视表
        pivot_table = df_filtered.pivot(
            index='months_since_purchase',
            columns='year',
            values=metric
        )
        # 绘制热力图并绑定到对应子图轴
        sns.heatmap(
            pivot_table,
            annot=True,
            cmap='YlBnBu',
            linewidth=2.5,
            linecolor='white',
            cbar=False,
            ax=axes[row, col]
        )
        # 设置子图标题与标签
        axes[row, col].set_title(f"{metric} - {company_size}", fontsize=14)
        axes[row, col].set_xlabel("Year", fontsize=12)
        axes[row, col].set_ylabel("Months Since First Purchase", fontsize=12)
        axes[row, col].tick_params(axis='both', which='major', labelsize=12)

# 自动调整子图间距，避免元素重叠
plt.tight_layout()
plt.show()

关键修正说明 #

• 手动指定size_order确保公司规模按要求排序
• 用axes[row, col]精准对应4×2布局的每个子图，避免变量覆盖
• 修正pivot的values参数为动态变量metric，正确读取指标数据
• 修复语法缩进与函数调用错误，保证代码可正常执行
• 添加plt.tight_layout()优化子图布局显示效果

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