用颜色映射的Matplotlib图例？

Question

总体上对matplotlib真的很困惑。我通常只使用导入matplotlib.pyplot作为plt。

然后做plt.figure（），plt.scatter（），plt.xlabel（），plt.show（）等等。但是我在谷歌上搜索如何做一些事情，比如用颜色映射图例，我得到了所有这些包括ax的例子。但是有plt.legend（），matplotlib留档中的例子只是显示plt.legend（句柄），但没有显示句柄应该是什么。如果我想做ax的事情，那么我必须重写我所有的代码，因为我想使用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”东西重写所有内容？如果没有，我该如何使其适应ax？我是否只写ax=plt.scatter（……）？

Answer 1

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

导入matplotlib。pyplot as 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()

输出: