Matplotlib 图例用颜色映射？

Question

一般来说，我真的对 matplotlib 感到困惑。 我通常只使用 import matplotlib.pyplot 作为 plt。

然后做所有的事情，比如 plt.figure()、plt.scatter()、plt.xlabel()、plt.show() 等等。但是我用谷歌搜索了如何做类似的事情，map 带有颜色的图例，我得到了所有这些例子包括 ax。 但是有 plt.legend() 和 matplotlib 文档中的示例只显示 plt.legend(handles) 但没有显示应该是什么句柄。 如果我想做斧头的事情，那么我必须重新编写我的所有代码，因为我想使用 plt，因为它更简单。

这是我的代码：

import matplotlib.pyplot as plt

colmap = {
    "domestic": "blue",
    "cheetah": "red",
    "leopard": "green",
    "tiger": "black"
}

colours = []

for i in y_train:
    colours.append(colmap[i])
    
plt.figure(figsize= [15,5])
plt.scatter(X_train[:,0], X_train[:,2],c=colours)
plt.xlabel('weight')
plt.ylabel('height')
plt.grid()
plt.show()

现在我想添加一个图例，它只显示与我的字典中相同的颜色。 但如果我这样做：

plt.legend(["国产","猎豹","豹子","老虎"])

它仅在图例中显示“国内”，颜色为红色，实际上与我对其进行颜色编码的方式不匹配。 有没有办法做到这一点而不用“斧头”的东西重写一切？ 如果没有，我该如何适应斧头？ 我只写 ax = plt.scatter(....) 吗？

Answer 1

未提供数据，但此代码可以帮助您了解如何在 matplotlib 中为散点图 plot 添加颜色：

导入 matplotlib.pyplot 作为 plt 导入 numpy 作为 np

# data for scatter plots
x = list(range(0,30))
y = [i**2 for i in x]

# data for mapping class to color
y_train = ['domestic','cheetah', 'cheetah', 'tiger', 'domestic',
           'leopard', 'tiger', 'domestic', 'cheetah', 'domestic',
           'leopard', 'leopard', 'domestic', 'domestic', 'domestic',
           'domestic', 'cheetah', 'tiger', 'cheetah', 'cheetah',
           'domestic', 'domestic', 'domestic', 'cheetah', 'leopard',
           'cheetah', 'domestic', 'cheetah', 'tiger', 'domestic']

# color mapper
colmap = {
    "domestic": "blue",
    "cheetah": "red",
    "leopard": "green",
    "tiger": "black"
}

# create color array
colors = [colmap[i] for i in y_train]

# plot scatter    
plt.figure(figsize=(15,5))
plt.scatter(x, y, c=colors)
plt.xlabel('weight')
plt.ylabel('height')
plt.grid()
plt.show()

Output：