python - Matplotlib - 向散点图添加图例

Question

我正在为每个人阅读熊猫书。在第 3 章中，作者使用以下代码创建了一个散点图：

# create a color variable based on sex
def recode_sex(sex):
    if sex == 'Female':
        return 0
    else:
        return 1

tips['sex_color'] = tips['sex'].apply(recode_sex)

scatter_plot = plt.figure(figsize=(20, 10))
axes1 = scatter_plot.add_subplot(1, 1, 1)
axes1.scatter(
    x=tips['total_bill'],
    y=tips['tip'],

    # set the size of the dots based on party size
    # we multiply the values by 10 to make the points bigger
#     and to emphasize the differences
    s=tips['size'] * 90,

#     set the color for the sex
    c=tips['sex_color'],

    # set the alpha value so points are more transparent
    # this helps with overlapping points
    alpha=0.5
)

axes1.set_title('Total Bill vs Tip Colored by Sex and Sized by Size')
axes1.set_xlabel('Total Bill')
axes1.set_ylabel('Tip')

plt.show()

情节如下所示：

我的问题是如何在散点图中添加图例？

标签： pythonmatplotlibscatter-plot

Answer 1

这是一个解决方案。此代码基于Matplotlib 的关于带有图例的散点图的教程。按性别分组的数据集的循环允许生成每个性别的颜色（和相应的图例）。然后从scatter函数的输出中指示大小，legend_elements用于大小。

这是我从您的示例中使用的数据集获得的：

这是代码：

import matplotlib.pyplot as plt
import seaborn as sns

# Read and group by gender
tips = sns.load_dataset("tips")
grouped = tips.groupby("sex")

# Show per group
fig, ax = plt.subplots(1)
for i, (name, group) in enumerate(grouped):
    sc = ax.scatter(
        group["total_bill"],
        group["tip"],
        s=group["size"] * 20,
        alpha=0.5,
        label=name,
    )

# Add legends (one for gender, other for size)
ax.add_artist(ax.legend(title='Gender'))
ax.legend(*sc.legend_elements("sizes", num=6), loc="lower left", title="Size")
ax.set_title("Scatter with legend")

plt.show()