绘制散点图，其中点的颜色取决于数据的类别

Question

我有一个数据列表X，它是一个Nx2矩阵。我想绘制每个X的第一元素相对于X的所有第二元素的图。首先，我已经将X的所有第一和第二元素分离到它们自己的列表中：X_comp1、X_comp2

我还有一个Nx1类别列表(cat(，它显示了X的元素属于哪个类别，即如果cat[i] = 3，则意味着X[i]属于类别3。

我想在散点图中为每个类别使用不同颜色的点。

到目前为止，我只能通过硬编码来实现这一点，但当类别更多时，这将变得非常低效。这是我的代码(假设有5个类别(：

#sample of X data
X = [[-0.13085379 -0.05958517],[ 0.02593188 -0.17576942],[-0.12505032 -0.02709171],[ 0.09790905 -0.18046944],[ 0.06437596 -0.20600157],[ 0.16287853 -0.2955353 ],[-0.52093842  0.33463338],[-0.03240038 -0.05431373],[-0.09645192 -0.14241157],[ 0.0807245  -0.26893815]]
X_comp1 = []#hold all the first components of X
X_comp2 = []#hold all the second components of X
cat = [1,3,2,1,5,3,2,4,4,1] 
#for testing just use 10 values, full file has over 3000 entries and 50 categories
for i in range(10):
X_comp1.append(X[i][0])
X_comp2.append(X[i][1])
for x1,x2,c in zip(X_comp1,X_comp2,cat):
if c == 1:
plt.scatter(x1,x2,c = 'b')
elif c == 2:
plt.scatter(x1,x2,c = 'g')
elif c == 3:
plt.scatter(x1,x2,c = 'r')
elif c == 4:
plt.scatter(x1,x2,c = 'c')
elif c == 5:
plt.scatter(x1,x2,c = 'm')
plt.legend([1,2,3,4,5])
plt.show()

我想让它对类别的数量更加灵活，这样我就不必为每个类别编写大量的if语句。

为了实现这一点，我想到了一个颜色列表：

colours = ["b", "g", "r", "c", "m",...]#number of colours depends on no. of categories
#This is the only element which I would like remain hard coded, as I want to choose the colours

其中每种颜色对应一个类别。然后，程序遍历所有数据，并根据类别绘制每个点。但我不确定这是如何实现的。

Answer 1

试试这个

color_dict = {1: 'b', 2: 'g', 3: 'r', 4: 'c', 5: 'm'}
for x1, x2, c in zip(X_comp1, X_comp2, cat):
if c in color_dict:
plt.scatter(x1, x2, c = color_dict[c])
plt.legend(list(color_dict.keys()))
plt.show()

我们可以删除所有if语句，而不是使用字典来检查c的每个值

Answer 2

对于漂亮的绘图，您也可以使用seaborn:

import seaborn as sns
import pandas as pd
sns.set()
df = pd.DataFrame({'X_1': X_comp1,'X_2':X_comp2, 'category':cat})
sns.scatterplot(data=df,x='X_1', y='X_2', hue='category')

如果你关心哪个类别应该有什么颜色，你可以通过palette参数和你自己的类别颜色dict:

my_palette = {1: 'b', 2: 'g', 3: 'r', 4: 'c', 5: 'm'}
sns.scatterplot(data=df,x='X_1', y='X_2', hue='category', palette=my_palette)

如果你对seaborn的默认选择不满意，还有一堆预定义的调色板。