Plotly子图：多色子图图例去重及添加State标题问题

问题场景 #

现有如下结构的DataFrame：

import pandas as pd

df = pd.DataFrame({
    "A": [1, 2, 4, 4],
    "B": [3, 10, 1, 3],
    "C": [5, 5, 7, 2],
    "D": [2, 1, 6, 8],
    "State": ["INITIAL", "DONE", "IN_PROGRESS", "PAUSED"]
})

需要创建垂直排列的两个散点子图（分别是A-B、A-C），期望实现的效果为：两个子图共享一套带有**"State"**标题的图例，每个状态对应唯一颜色，无重复图例条目。

用户尝试的代码如下，但存在图例重复（每个子图都生成一套State图例）、缺少图例标题的问题：

import plotly.graph_objects as go
from plotly.subplots import make_subplots

fig = go.Figure()

figures = [
        df.plot.scatter(x="A", y="B", color="State"),
        df.plot.scatter(x="A", y="C", color="State")
    ]

fig = make_subplots(rows=len(figures), cols=1) 

for i, figure in enumerate(figures):
    for trace in range(len(figure["data"])):
        fig.append_trace(figure["data"][trace], row=i+1, col=1)

fig.update_xaxes(title_text="A", row=2, col=1)  
fig.update_xaxes(title_text="A", row=1, col=1)       
fig.update_yaxes(title_text="C", row=2, col=1)
fig.update_yaxes(title_text="B", row=1, col=1)
fig.show()

解决方案 #

核心思路：

• 仅保留每个State对应的第一个图例条目，其余子图中同State的trace隐藏图例
• 通过update_layout添加图例标题
• 统一颜色映射，确保同一State在两个子图中颜色一致

修改后的完整代码：

import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pandas as pd

df = pd.DataFrame({
    "A": [1, 2, 4, 4],
    "B": [3, 10, 1, 3],
    "C": [5, 5, 7, 2],
    "D": [2, 1, 6, 8],
    "State": ["INITIAL", "DONE", "IN_PROGRESS", "PAUSED"]
})

# 复用Plotly默认颜色序列，生成State与颜色的映射
state_list = df["State"].unique()
color_map = {state: go.scatter.marker.colorway[i] for i, state in enumerate(state_list)}

fig = make_subplots(rows=2, cols=1)

# 记录已添加图例的State，避免重复
added_legends = set()

# 绘制第一个子图：A vs B
for state in state_list:
    state_data = df[df["State"] == state]
    show_legend = state not in added_legends
    fig.add_trace(
        go.Scatter(
            x=state_data["A"],
            y=state_data["B"],
            mode="markers",
            name=state,
            marker=dict(color=color_map[state]),
            showlegend=show_legend
        ),
        row=1, col=1
    )
    if show_legend:
        added_legends.add(state)

# 绘制第二个子图：A vs C（隐藏重复图例）
for state in state_list:
    state_data = df[df["State"] == state]
    fig.add_trace(
        go.Scatter(
            x=state_data["A"],
            y=state_data["C"],
            mode="markers",
            name=state,
            marker=dict(color=color_map[state]),
            showlegend=False
        ),
        row=2, col=1
    )

# 设置坐标轴标题
fig.update_xaxes(title_text="A", row=1, col=1)
fig.update_xaxes(title_text="A", row=2, col=1)
fig.update_yaxes(title_text="B", row=1, col=1)
fig.update_yaxes(title_text="C", row=2, col=1)

# 添加图例标题并调整布局
fig.update_layout(
    legend_title_text="State",
    height=600,
    width=800
)

fig.show()

关键修改说明 #

• 避免图例重复：通过added_legends集合跟踪已添加图例的State，第一个子图中首次出现的State显示图例，第二个子图中所有State的trace设置showlegend=False
• 添加图例标题：使用fig.update_layout(legend_title_text="State")直接指定图例标题
• 统一颜色：复用Plotly默认颜色序列，确保同一State在两个子图中颜色完全一致

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