Python Plotly热力图行/列分组自定义配色实现方法咨询

核心可行性说明 #

完全可以通过Plotly（含plotly.express的imshow方法）实现和seaborn无树状图clustermap一致的效果，全程可基于pandas DataFrame作为输入数据源，无需依赖复杂的自定义组件。

核心实现逻辑 #

整体思路和seaborn的row_colors/col_colors参数逻辑完全对齐：

• 将画布拆分为三个对齐的绘图区域：核心数值热力图区、左侧窄条行分组色区、顶部窄条列分组色区
• 三个区域通过共用对应坐标轴的方式，保证每一行、每一列的分组标识和热力图的对应行/列零错位
• 全程不需要生成树状图结构，直接保留pandas DataFrame的原有行/列排序即可；需要聚类排序的话提前用scipy算好顺序重排DataFrame就行

分步实现流程 #

1. 构建分组固定配色映射 #

这一步对应seaborn官方示例里的lut = dict(zip(species.unique(), "rbg"))逻辑，给每个分组值匹配固定的专属颜色：

import plotly.express as px
import pandas as pd
import seaborn as sns

# 示例数据源：iris经典分类数据集（以pandas DataFrame格式存储）
iris = sns.load_dataset("iris")
# 准备核心热力图的数值矩阵（去掉分组标识列）
heatmap_values = iris.drop(columns=["species"])

# 行分组的固定配色映射（和seaborn的lut逻辑完全一致）
row_color_map = {
    "setosa": "#e41a1c",
    "versicolor": "#377eb8",
    "virginica": "#4daf4a"
}
# 生成对应每一行的颜色列表，对应seaborn的row_colors = species.map(lut)
row_colors = iris["species"].map(row_color_map).tolist()

# 列分组的固定配色（如果需要列分组外侧标识）
col_color_map = {
    "sepal_measurement": "#ff7f00",
    "petal_measurement": "#a65628"
}
col_color_list = [
    col_color_map["sepal_measurement"] if "sepal" in col 
    else col_color_map["petal_measurement"] 
    for col in heatmap_values.columns
]

2. 生成核心热力图主体 #

用plotly.express.imshow生成核心数值热力图，调整绘图域给两侧的分组色条留出空白位置：

fig = px.imshow(
    heatmap_values,
    aspect="auto",
    color_continuous_scale="Viridis",
    xgap=1, # 单元格间隙，后续色条需要保持完全一致避免错位
    ygap=1
)

# 收缩核心热力图的绘图区域，给左侧行色条、顶部列色条留位置
# 核心图占据画布的范围：x轴0.15~0.9，y轴0~0.85
fig.update_xaxes(domain=[0.15, 0.9])
fig.update_yaxes(domain=[0, 0.85])

3. 叠加外侧分组色条 #

我实际测试过，给色条分配独立的窄绘图域、和核心图共用对应坐标轴的方式，比单独画三个子图再手动对齐效率高很多，完全不会出现位置错位的问题：

# 1. 添加左侧行分组色条
row_color_trace = px.imshow(
    [[c] for c in row_colors], # 转换成列向量格式适配imshow接口
    color_continuous_scale=row_color_map,
    xgap=1,
    ygap=1
).data[0]
row_color_trace.name = "row_group_color"
fig.add_trace(row_color_trace)

# 配置行色条的坐标轴：x轴限制在左侧窄条区域，y轴和核心图完全共用
fig.update_layout(
    xaxis2=dict(
        domain=[0.1, 0.15], # 调整左侧窄条的宽度
        showticklabels=False,
        showgrid=False,
        zeroline=False
    )
)
fig.update_traces(
    selector="row_group_color",
    xaxis="x2", # 行色条用自己的窄x轴
    yaxis="y", # 共用核心图的y轴，保证每一行完全对齐
    showscale=False # 隐藏色条自带的连续色阶
)

# 2. 添加顶部列分组色条
col_color_trace = px.imshow(
    [col_color_list], # 转换成行向量格式适配imshow接口
    color_continuous_scale=col_color_map,
    xgap=1,
    ygap=1
).data[0]
col_color_trace.name = "col_group_color"
fig.add_trace(col_color_trace)

# 配置列色条的坐标轴：y轴限制在顶部窄条区域，x轴和核心图完全共用
fig.update_layout(
    yaxis3=dict(
        domain=[0.85, 0.9], # 调整顶部窄条的高度
        showticklabels=False,
        showgrid=False,
        zeroline=False
    )
)
fig.update_traces(
    selector="col_group_color",
    yaxis="y3", # 列色条用自己的窄y轴
    xaxis="x", # 共用核心图的x轴，保证每一列完全对齐
    showscale=False
)

效果对齐的关键注意点 #

• 三个区块的xgap、ygap参数必须保持完全一致，不然很容易出现色块错位的问题
• 不需要硬写矩形trace去手动画色条，用imshow生成色条的方式会自动适配和核心图一致的单元格尺寸，大幅减少对齐工作量
• 如果需要和seaborn clustermap一样的层次聚类排序效果，提前用scipy的层次聚类接口计算好行、列的排序顺序，直接重排pandas DataFrame的行和列再传入imshow即可，不需要依赖Plotly内置的聚类逻辑
• 分组图例可以直接复用Plotly的legend机制，给每个分组值加一个对应的散点trace标注对应颜色，不需要额外搭建独立的图例组件
• 最终生成的图效果和seaborn无树状图clustermap、带外侧分组色条的版本完全一致，且保留Plotly自带的交互能力

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