基于时长创建单元格宽度可变的Seaborn热力图

问题背景 #

需要让热力图的每个单元格宽度与DataFrame中Duration字段的事件时长成正比，当前示例包含3个事件，默认Seaborn热力图为等宽单元格，无法满足需求。

解决方案 #

Seaborn原生heatmap不支持不等宽单元格，可通过Matplotlib手动绘制矩形实现，步骤如下：

• 解析时长字段为可计算的数值（总天数）
• 计算每个单元格的相对宽度比例
• 逐个绘制对应宽度、颜色的矩形，并添加标签
• 配置颜色条与坐标轴样式

完整代码 #

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import pandas as pd

# 初始化数据
data = pd.DataFrame(
    {
        'Date': ['2021-01-07', '2021-06-07', '2021-07-21'],
        'Duration': ['5 days 03:00:00', '15 days 02:00:00', '2 days 10:00:00'],
        'L': [3, 1, 3],
        'Pe': [1, 3, 1]
    }
)
data.set_index('Date', inplace=True)

# 将Duration转换为总天数（保留小数）
data['Duration_days'] = pd.to_timedelta(data['Duration']).dt.total_seconds() / (24 * 3600)

# 提取绘图所需数据
row_labels = ['L', 'Pe']
col_labels = data.index.tolist()
value_matrix = data[['L', 'Pe']].values.T

# 配置颜色映射
cmap = matplotlib.cm.get_cmap('viridis')
color_list = [cmap(0.25), cmap(0.5), cmap(0.75), cmap(1)]
# 匹配数值与颜色（1-3对应前3种颜色）
value_color_map = {1: color_list[0], 2: color_list[1], 3: color_list[2]}

# 计算每个列的相对宽度
total_duration = data['Duration_days'].sum()
col_widths = data['Duration_days'] / total_duration

# 创建画布
fig, ax = plt.subplots(figsize=(15, 3))

# 遍历绘制每个单元格
y_coord = 0
for row_name in row_labels:
    x_coord = 0
    for idx, col_name in enumerate(col_labels):
        # 获取当前单元格的宽度和颜色
        cell_width = col_widths[idx]
        cell_value = value_matrix[y_coord][idx]
        cell_color = value_color_map[cell_value]
        
        # 绘制矩形单元格
        rect = patches.Rectangle((x_coord, y_coord), cell_width, 1, color=cell_color)
        ax.add_patch(rect)
        
        # 添加列标签（居中）
        ax.text(
            x_coord + cell_width / 2,
            y_coord + 0.5,
            col_name,
            ha='center',
            va='center',
            fontsize=10
        )
        x_coord += cell_width
    
    # 添加行标签
    ax.text(
        -0.02,
        y_coord + 0.5,
        row_name,
        ha='right',
        va='center',
        fontsize=12
    )
    y_coord += 1

# 设置坐标轴范围与样式
ax.set_xlim(0, 1)
ax.set_ylim(0, len(row_labels))
ax.set_xticks([])
ax.set_yticks([])

# 添加颜色条
sm = matplotlib.cm.ScalarMappable(
    cmap=matplotlib.colors.ListedColormap(color_list[:3])
)
sm.set_array([1, 2, 3])
cbar = plt.colorbar(sm, ax=ax, ticks=[1, 2, 3])
cbar.set_label('Classes')

plt.tight_layout()
plt.show()

关键说明 #

1. 时长解析：利用pd.to_timedelta将字符串格式的时长转换为时间差，再计算总天数，为宽度比例计算提供数值基础。
2. 宽度比例：通过每个事件时长占总时长的比例，确保单元格宽度与时长成正比，总宽度适配坐标轴0-1的范围。
3. 矩形绘制：使用matplotlib.patches.Rectangle逐个绘制单元格，同时添加行/列标签保证可读性。
4. 颜色适配：根据数据中的类别值（1-3）匹配对应的颜色，并生成对应的颜色条。

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