基于groupby为VAR1分组生成子图并绘制VAR2散点图的实现问题

给定如下数据集与初始化代码：

data = {0: {'VAR1': 'A', 'VAR2': 'X', 'VAL1': 3, 'VAL2': 1},
 1: {'VAR1': 'A', 'VAR2': 'X', 'VAL1': 4, 'VAL2': 1},
 2: {'VAR1': 'A', 'VAR2': 'X', 'VAL1': 5, 'VAL2': 1},
 3: {'VAR1': 'A', 'VAR2': 'Y', 'VAL1': 3, 'VAL2': 2},
 4: {'VAR1': 'A', 'VAR2': 'Y', 'VAL1': 4, 'VAL2': 2},
 5: {'VAR1': 'A', 'VAR2': 'Y', 'VAL1': 5, 'VAL2': 2},
 6: {'VAR1': 'A', 'VAR2': 'Z', 'VAL1': 3, 'VAL2': 3},
 7: {'VAR1': 'A', 'VAR2': 'Z', 'VAL1': 4, 'VAL2': 3},
 8: {'VAR1': 'A', 'VAR2': 'Z', 'VAL1': 5, 'VAL2': 3},
 9: {'VAR1': 'B', 'VAR2': 'X', 'VAL1': 3, 'VAL2': 1},
 10: {'VAR1': 'B', 'VAR2': 'X', 'VAL1': 4, 'VAL2': 1},
 11: {'VAR1': 'B', 'VAR2': 'X', 'VAL1': 5, 'VAL2': 1},
 12: {'VAR1': 'B', 'VAR2': 'Y', 'VAL1': 3, 'VAL2': 2},
 13: {'VAR1': 'B', 'VAR2': 'Y', 'VAL1': 4, 'VAL2': 2},
 14: {'VAR1': 'B', 'VAR2': 'Y', 'VAL1': 5, 'VAL2': 2},
 15: {'VAR1': 'B', 'VAR2': 'Z', 'VAL1': 3, 'VAL2': 3},
 16: {'VAR1': 'B', 'VAR2': 'Z', 'VAL1': 4, 'VAL2': 3},
 17: {'VAR1': 'B', 'VAR2': 'Z', 'VAL1': 5, 'VAL2': 3},
 18: {'VAR1': 'C', 'VAR2': 'X', 'VAL1': 3, 'VAL2': 1},
 19: {'VAR1': 'C', 'VAR2': 'X', 'VAL1': 4, 'VAL2': 1},
 20: {'VAR1': 'C', 'VAR2': 'X', 'VAL1': 5, 'VAL2': 1},
 21: {'VAR1': 'C', 'VAR2': 'Y', 'VAL1': 3, 'VAL2': 2},
 22: {'VAR1': 'C', 'VAR2': 'Y', 'VAL1': 4, 'VAL2': 2},
 23: {'VAR1': 'C', 'VAR2': 'Y', 'VAL1': 5, 'VAL2': 2},
 24: {'VAR1': 'C', 'VAR2': 'Z', 'VAL1': 3, 'VAL2': 3},
 25: {'VAR1': 'C', 'VAR2': 'Z', 'VAL1': 4, 'VAL2': 3},
 26: {'VAR1': 'C', 'VAR2': 'Z', 'VAL1': 5, 'VAL2': 3}}

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame.from_dict(data, orient='index')

需求 #

• 为VAR1中的每个唯一值创建独立坐标轴
• 在每个坐标轴中，针对VAR2的不同分组绘制VAL1（x轴）与VAL2（y轴）的散点图

尝试过的方案 #

最初的方案存在逻辑错误：

group_var1 = df.groupby('VAR1')

for name_var1, grouped_var1 in group_var1:
    i = 0
    fig, axes = plt.subplots(nrows=3, ncols=1,figsize=(20, 8), tight_layout=True)
    group_var2 = grouped_var1.groupby('VAR2')
    for name_var2, grouped_var2 in group_var2:
        grouped_var2.plot(kind='scatter', ax=axes[i], x='VAL1', y='VAL2')

    i+=1

修改后的方案可运行，但实现不够优雅：

group_var1 = df.groupby('VAR1')
fig, axes = plt.subplots(nrows=3, ncols=1,figsize=(20, 8), tight_layout=True)
i = 0
for name_var1, grouped_var1 in group_var1:
    group_var2 = grouped_var1.groupby('VAR2')
    for name_var2, grouped_var2 in group_var2:
        grouped_var2.plot(kind='scatter', ax=axes[i], x='VAL2', y='VAL1', c=['red','green','yellow'])
    i+=1

寻求更优、更优雅的基于groupby的实现方式。

以下是优化后的代码，逻辑清晰且易于维护：

import pandas as pd
import matplotlib.pyplot as plt

# 初始化数据（同上，可保留）
data = {0: {'VAR1': 'A', 'VAR2': 'X', 'VAL1': 3, 'VAL2': 1},
 1: {'VAR1': 'A', 'VAR2': 'X', 'VAL1': 4, 'VAL2': 1},
 2: {'VAR1': 'A', 'VAR2': 'X', 'VAL1': 5, 'VAL2': 1},
 3: {'VAR1': 'A', 'VAR2': 'Y', 'VAL1': 3, 'VAL2': 2},
 4: {'VAR1': 'A', 'VAR2': 'Y', 'VAL1': 4, 'VAL2': 2},
 5: {'VAR1': 'A', 'VAR2': 'Y', 'VAL1': 5, 'VAL2': 2},
 6: {'VAR1': 'A', 'VAR2': 'Z', 'VAL1': 3, 'VAL2': 3},
 7: {'VAR1': 'A', 'VAR2': 'Z', 'VAL1': 4, 'VAL2': 3},
 8: {'VAR1': 'A', 'VAR2': 'Z', 'VAL1': 5, 'VAL2': 3},
 9: {'VAR1': 'B', 'VAR2': 'X', 'VAL1': 3, 'VAL2': 1},
 10: {'VAR1': 'B', 'VAR2': 'X', 'VAL1': 4, 'VAL2': 1},
 11: {'VAR1': 'B', 'VAR2': 'X', 'VAL1': 5, 'VAL2': 1},
 12: {'VAR1': 'B', 'VAR2': 'Y', 'VAL1': 3, 'VAL2': 2},
 13: {'VAR1': 'B', 'VAR2': 'Y', 'VAL1': 4, 'VAL2': 2},
 14: {'VAR1': 'B', 'VAR2': 'Y', 'VAL1': 5, 'VAL2': 2},
 15: {'VAR1': 'B', 'VAR2': 'Z', 'VAL1': 3, 'VAL2': 3},
 16: {'VAR1': 'B', 'VAR2': 'Z', 'VAL1': 4, 'VAL2': 3},
 17: {'VAR1': 'B', 'VAR2': 'Z', 'VAL1': 5, 'VAL2': 3},
 18: {'VAR1': 'C', 'VAR2': 'X', 'VAL1': 3, 'VAL2': 1},
 19: {'VAR1': 'C', 'VAR2': 'X', 'VAL1': 4, 'VAL2': 1},
 20: {'VAR1': 'C', 'VAR2': 'X', 'VAL1': 5, 'VAL2': 1},
 21: {'VAR1': 'C', 'VAR2': 'Y', 'VAL1': 3, 'VAL2': 2},
 22: {'VAR1': 'C', 'VAR2': 'Y', 'VAL1': 4, 'VAL2': 2},
 23: {'VAR1': 'C', 'VAR2': 'Y', 'VAL1': 5, 'VAL2': 2},
 24: {'VAR1': 'C', 'VAR2': 'Z', 'VAL1': 3, 'VAL2': 3},
 25: {'VAR1': 'C', 'VAR2': 'Z', 'VAL1': 4, 'VAL2': 3},
 26: {'VAR1': 'C', 'VAR2': 'Z', 'VAL1': 5, 'VAL2': 3}}

df = pd.DataFrame.from_dict(data, orient='index')

# 定义VAR2分组的统一颜色映射
color_map = {'X': 'red', 'Y': 'green', 'Z': 'yellow'}
# 动态获取VAR1唯一值数量，创建对应数量的子图
n_var1_groups = df['VAR1'].nunique()
fig, axes = plt.subplots(nrows=n_var1_groups, ncols=1, figsize=(20, 8), tight_layout=True)

# 遍历VAR1分组，用enumerate直接获取子图索引
for ax_idx, (var1_val, var1_group) in enumerate(df.groupby('VAR1')):
    current_ax = axes[ax_idx]
    # 遍历当前VAR1分组下的VAR2子组
    for var2_val, var2_group in var1_group.groupby('VAR2'):
        var2_group.plot(
            kind='scatter',
            x='VAL1',
            y='VAL2',
            color=color_map[var2_val],
            label=f'VAR2={var2_val}',
            ax=current_ax
        )
    # 设置子图标题与坐标轴标签，提升可读性
    current_ax.set_title(f'VAR1 = {var1_val}')
    current_ax.set_xlabel('VAL1')
    current_ax.set_ylabel('VAL2')
    current_ax.legend()

plt.show()

方案优势 #

• 动态适配数据：通过df['VAR1'].nunique()自动确定子图数量，无需硬编码行数，适配不同规模的数据集
• 统一视觉语言：用color_map保证同一VAR2分组在所有子图中颜色一致，提升可视化一致性
• 简洁的索引管理：使用enumerate直接绑定子图索引，避免手动维护计数器带来的逻辑错误
• 信息直观清晰：添加子图标题、坐标轴标签和图例，让图表传递的信息更明确
• 逻辑分层清晰：外层遍历VAR1分组，内层遍历VAR2分组，代码结构一目了然，便于后续维护

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