如何用Python绘制含分类变量的3D柱状图并处理负数值？

需求说明 #

需基于包含3个数值变量的数据集绘制3D柱状图，核心要求：

• 用pd.cut()将Fibrinogen和RR__root_mean_square划分为4个区间的分类变量
• 处理体重增长（Weight gain）的负值：通过颜色区分正负值
• x、y轴从零开始，优化图表可读性

数据集示例 #

Fibrinogen  RR__root_mean_square  Weight gain
2         5.15             26.255383         -2.0
3         0.99             20.934106          0.5
7         2.12             24.252434         -1.0
11        1.64             17.004289          3.0
12        3.06             21.201716          0.0

原代码问题 #

原代码使用cat.codes转换分类变量导致丢失区间标签，且3D柱状图参数设置不当，图表可读性差。原代码如下：

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# sample data
data = {'Fibrinogen': [5.15, 0.99, 2.12, 1.64, 3.06],
        'RR__root_mean_square': [26.255383, 20.934106, 24.252434, 17.004289, 21.201716],
        'Weight gain': [-2.0, 0.5, -1.0, 3.0, 0.0]}

X_success = pd.DataFrame(data)

x1 = pd.cut(X_success['Fibrinogen'], [0,2,4,10]).cat.codes
y1 = pd.cut(X_success['RR__root_mean_square'], [15,25,32.5,40]).cat.codes
z1 = X_success['Weight gain']

# cat.codes to don't have the 'float() argument must be a string or a number, not 'pandas._libs.interval.Interval'' error
# But I loose the label...

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

ax.bar(x1, y1, z1, zdir='y', alpha=0.8)
plt.show()

优化后代码 #

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# 样本数据
data = {'Fibrinogen': [5.15, 0.99, 2.12, 1.64, 3.06],
        'RR__root_mean_square': [26.255383, 20.934106, 24.252434, 17.004289, 21.201716],
        'Weight gain': [-2.0, 0.5, -1.0, 3.0, 0.0]}
df = pd.DataFrame(data)

# 用pd.cut生成分类变量并保留区间标签
fibrinogen_bins = [0, 2, 4, 10]
fibrinogen_labels = ['0-2', '2-4', '4-10']
df['Fibrinogen_cat'] = pd.cut(df['Fibrinogen'], bins=fibrinogen_bins, labels=fibrinogen_labels)
x_codes = df['Fibrinogen_cat'].cat.codes

rr_bins = [15, 25, 32.5, 40]
rr_labels = ['15-25', '25-32.5', '32.5-40']
df['RR_cat'] = pd.cut(df['RR__root_mean_square'], bins=rr_bins, labels=rr_labels)
y_codes = df['RR_cat'].cat.codes

weight_gain = df['Weight gain']

# 区分正负值的颜色
colors = ['#ff6b6b' if val < 0 else '#4ecdc4' for val in weight_gain]

# 创建3D图表
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 设置柱子宽度
dx = dy = 0.5

# 绘制每个柱子，负值从0向下延伸，正值向上
for x, y, z, color in zip(x_codes, y_codes, weight_gain, colors):
    if z >= 0:
        ax.bar3d(x, y, 0, dx, dy, z, color=color, alpha=0.8)
    else:
        ax.bar3d(x, y, z, dx, dy, -z, color=color, alpha=0.8)

# 设置轴标签和刻度
ax.set_xlabel('Fibrinogen')
ax.set_xticks(np.arange(len(fibrinogen_labels)))
ax.set_xticklabels(fibrinogen_labels)

ax.set_ylabel('RR Root Mean Square')
ax.set_yticks(np.arange(len(rr_labels)))
ax.set_yticklabels(rr_labels)

ax.set_zlabel('Weight Gain')

# 设置x、y轴从零开始
ax.set_xlim(-0.2, len(fibrinogen_labels)-0.8)
ax.set_ylim(-0.2, len(rr_labels)-0.8)

plt.tight_layout()
plt.show()

关键优化点 #

• 保留分类标签：给pd.cut()添加labels参数定义区间名称，同时用cat.codes获取绘图坐标值，最后将标签绑定到轴刻度，避免信息丢失
• 负值颜色区分：通过列表推导式根据体重增长值生成颜色列表，负值用红色，正值用蓝绿色，直观区分增减重
• 轴归零处理：调整xlim和ylim让轴起点贴近0，优化布局紧凑度
• 3D柱状图逻辑修正：针对负值柱子，从z值位置向上绘制长度为-z的柱子，确保柱子从0平面向上下两个方向延伸，符合数据含义

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