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()

剧情如下：

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

Answer 1

这里有一个解决方案。此代码基于Matplotlib's tutorial on scatter plot with legends。按性别分组的数据集的循环允许生成每个性别的颜色（和相应的图例）。然后从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()