如何根据Department参数设置散点图气泡颜色

问题描述 #

现有如下使用Plotly Express绘制散点图的Python代码，希望根据Department字段的值为不同部门的气泡设置专属颜色：

#the data
data = pd.DataFrame({
    'label': ['Luis', 'Sara', 'Jeroen', 'Sophie', 'Florence', 'Simeon', 'Lambert', 'Rahul', 'Peter', 'John'],
    'Data Culture': [0, 5, 10, 15, 20, 18, 13, 9, 4, 1],
    'Data Skills': [20, 15, 10, 5, 0, 3, 7, 11, 16, 18],
    'Department' :['Dept1', 'Dept2', 'Dept3', 'Dept2', 'Dept1', 'Dept2', 'Dept1', 'Dept2', 'Dept1', 'Dept2']
    })
print(data)

fig = px.scatter(data, x=data["Data Culture"], y=data["Data Skills"], text=data.label, 
                 title='Data Culture vs Data Skills',
                 width=800, height=600)

# calculate averages
x_avg = data['Data Culture'].mean()
y_avg = data['Data Skills'].mean()

# add horizontal and vertical lines
fig.add_vline(x=10, line_width=3, opacity=0.5)
fig.add_hline(y=10, line_width=3, opacity=0.5)

# set x limits
adj_x = max((data['Data Culture'].max() - x_avg), (x_avg - data['Data Culture'].min())) * 1.1
lb_x, ub_x = (x_avg - adj_x, x_avg + adj_x)
fig.update_xaxes(range = [lb_x, ub_x])

# set y limits
adj_y = max((data['Data Skills'].max() - y_avg), (y_avg - data['Data Skills'].min())) * 1.1
lb_y, ub_y = (y_avg - adj_y, y_avg + adj_y)
fig.update_yaxes(range = [lb_y, ub_y])

# update x tick labels
axis = ['Low', 'High']     
fig.update_layout(
    xaxis_title='Data Culture',
    xaxis = dict(
        tickmode = 'array',
        tickvals = ([(x_avg - adj_x / 2), (x_avg + adj_x / 2)]),
        ticktext = axis
      )
    )

# update y tick labels
fig.update_layout(
    yaxis_title='Data Skills',
    yaxis = dict(
        tickmode = 'array',
        tickvals = ([(y_avg - adj_y / 2), (y_avg + adj_y / 2)]),
        ticktext = axis,
        tickangle=270
        )
    ) 

fig.update_layout(margin=dict(t=50, l=5, r=5, b=50),
    title={'text': 'pl',
           'font_size': 20,
           'y':1.0,
           'x':0.5,
           'xanchor': 'center',
           'yanchor': 'top'})

# where I need the help with annotation
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=0, y=-0.15,#data['score'].min()-0.2, y=data['wgt'].min()-0.2,
                        text="Han Solo",
                        xref='paper',
                        yref='paper', 
                        showarrow=False))
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=1, y=-0.15,#x=data['score'].max(), y=data['wgt'].min(),
                        text="Young Padawan",
                        xref='paper',
                        yref='paper',
                        showarrow=False))
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=0, y=1.15, #x=data['score'].min(), y=data['wgt'].max(),
                        text="Baby Yoda",
                        xref='paper',
                        yref='paper',
                        showarrow=False))
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=1, y=1.15, #x=data['score'].max(), y=data['wgt'].max(),
                        text="Yoda",
                        xref='paper',
                        yref='paper',
                        showarrow=False))

fig.update_layout(
    margin=dict(l=20, r=20, t=100, b=100),
)


fig.show()

当前生成的散点图效果：

解决方案 #

只需在px.scatter()函数中添加color='Department'参数，即可让Plotly自动按部门为气泡分配不同颜色。如果需要自定义各部门的专属颜色，可以额外添加color_discrete_map参数指定颜色映射。

修改后的核心代码部分：

# 自定义颜色映射（可选）
custom_colors = {
    'Dept1': '#FF5733',
    'Dept2': '#33FF57',
    'Dept3': '#3357FF'
}

fig = px.scatter(data, x="Data Culture", y="Data Skills", text="label", 
                 title='Data Culture vs Data Skills',
                 width=800, height=600,
                 color='Department',  # 按部门设置颜色
                 color_discrete_map=custom_colors)  # 自定义颜色（可选）

完整修改后的代码 #

import pandas as pd
import plotly.express as px

#the data
data = pd.DataFrame({
    'label': ['Luis', 'Sara', 'Jeroen', 'Sophie', 'Florence', 'Simeon', 'Lambert', 'Rahul', 'Peter', 'John'],
    'Data Culture': [0, 5, 10, 15, 20, 18, 13, 9, 4, 1],
    'Data Skills': [20, 15, 10, 5, 0, 3, 7, 11, 16, 18],
    'Department' :['Dept1', 'Dept2', 'Dept3', 'Dept2', 'Dept1', 'Dept2', 'Dept1', 'Dept2', 'Dept1', 'Dept2']
    })
print(data)

# 自定义颜色映射（可选）
custom_colors = {
    'Dept1': '#FF5733',
    'Dept2': '#33FF57',
    'Dept3': '#3357FF'
}

fig = px.scatter(data, x="Data Culture", y="Data Skills", text="label", 
                 title='Data Culture vs Data Skills',
                 width=800, height=600,
                 color='Department',  # 按部门区分颜色
                 color_discrete_map=custom_colors)  # 自定义各部门颜色（可选）

# calculate averages
x_avg = data['Data Culture'].mean()
y_avg = data['Data Skills'].mean()

# add horizontal and vertical lines
fig.add_vline(x=10, line_width=3, opacity=0.5)
fig.add_hline(y=10, line_width=3, opacity=0.5)

# set x limits
adj_x = max((data['Data Culture'].max() - x_avg), (x_avg - data['Data Culture'].min())) * 1.1
lb_x, ub_x = (x_avg - adj_x, x_avg + adj_x)
fig.update_xaxes(range = [lb_x, ub_x])

# set y limits
adj_y = max((data['Data Skills'].max() - y_avg), (y_avg - data['Data Skills'].min())) * 1.1
lb_y, ub_y = (y_avg - adj_y, y_avg + adj_y)
fig.update_yaxes(range = [lb_y, ub_y])

# update x tick labels
axis = ['Low', 'High']     
fig.update_layout(
    xaxis_title='Data Culture',
    xaxis = dict(
        tickmode = 'array',
        tickvals = ([(x_avg - adj_x / 2), (x_avg + adj_x / 2)]),
        ticktext = axis
      )
    )

# update y tick labels
fig.update_layout(
    yaxis_title='Data Skills',
    yaxis = dict(
        tickmode = 'array',
        tickvals = ([(y_avg - adj_y / 2), (y_avg + adj_y / 2)]),
        ticktext = axis,
        tickangle=270
        )
    ) 

fig.update_layout(margin=dict(t=50, l=5, r=5, b=50),
    title={'text': 'pl',
           'font_size': 20,
           'y':1.0,
           'x':0.5,
           'xanchor': 'center',
           'yanchor': 'top'})

# where I need the help with annotation
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=0, y=-0.15,#data['score'].min()-0.2, y=data['wgt'].min()-0.2,
                        text="Han Solo",
                        xref='paper',
                        yref='paper', 
                        showarrow=False))
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=1, y=-0.15,#x=data['score'].max(), y=data['wgt'].min(),
                        text="Young Padawan",
                        xref='paper',
                        yref='paper',
                        showarrow=False))
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=0, y=1.15, #x=data['score'].min(), y=data['wgt'].max(),
                        text="Baby Yoda",
                        xref='paper',
                        yref='paper',
                        showarrow=False))
fig.add_annotation(dict(font=dict(color="black",size=18),
                        x=1, y=1.15, #x=data['score'].max(), y=data['wgt'].max(),
                        text="Yoda",
                        xref='paper',
                        yref='paper',
                        showarrow=False))

fig.update_layout(
    margin=dict(l=20, r=20, t=100, b=100),
)


fig.show()

说明 #

• 添加color='Department'后，Plotly会自动为每个部门分配不同颜色，并生成对应的图例。
• 如果不满意默认颜色，使用color_discrete_map参数可以手动指定每个部门的颜色，值为字典类型，键是部门名称，值是颜色代码（支持十六进制、RGB等格式）。

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