python - matplotlib 中带有图例和随机点顺序的散点图

Question

我正在尝试构建来自 python/matplotlib 中多个类的大量数据的散点图。不幸的是，我似乎必须在数据随机化和图例标签之间做出选择。有没有办法可以同时拥有两者(最好不需要手动编码标签？)

最小可重现示例:

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
X = np.random.normal(0, 1, [5000, 2])
Y = np.random.normal(0.5, 1, [5000, 2]) 
data = np.concatenate([X,Y])
classes = np.concatenate([np.repeat('X', X.shape[0]),
                          np.repeat('Y', Y.shape[0])])

用随机点绘图:

plot_idx = np.random.permutation(data.shape[0])
colors = pd.factorize(classes)
fig, ax = plt.subplots()
ax.scatter(data[plot_idx, 0], 
           data[plot_idx, 1], 
           c=colors[plot_idx],
           label=classes[plot_idx],
           alpha=0.4)
plt.legend()
plt.show()

这给了我错误的图例。

使用正确的图例进行绘图:

from matplotlib import cm
unique_classes = np.unique(classes)
colors = cm.Set1(np.linspace(0, 1, len(unique_classes)))
for i, class in enumerate(unique_classes):
    ax.scatter(data[classes == class, 0], 
               data[classes == class, 1],
               c=colors[i],
               label=class,
               alpha=0.4)
plt.legend()
plt.show()

但现在这些点不是随机的，结果图不能代表数据。

我正在寻找能给我一个结果的东西，就像我在 R 中得到的结果:

library(ggplot2)
X <- matrix(rnorm(10000, 0, 1), ncol=2)
Y <- matrix(rnorm(10000, 0.5, 1), ncol=2)
data <- as.data.frame(rbind(X, Y))
data$classes <- rep(c('X', 'Y'), times=nrow(X))
plot_idx <- sample(nrow(data))

ggplot(data[plot_idx,], aes(x=V1, y=V2, color=classes)) +
  geom_point(alpha=0.4, size=3)

Answer 1

您需要手动创建图例。但这并不是一个大问题。您可以循环遍历标签并为每个标签创建一个图例条目。这里可以使用带有类似于散点图的标记的 Line2D 作为句柄。

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
X = np.random.normal(0, 1, [5000, 2])
Y = np.random.normal(0.5, 1, [5000, 2]) 
data = np.concatenate([X,Y])
classes = np.concatenate([np.repeat('X', X.shape[0]),
                          np.repeat('Y', Y.shape[0])])

plot_idx = np.random.permutation(data.shape[0])
colors,labels = pd.factorize(classes)

fig, ax = plt.subplots()
sc = ax.scatter(data[plot_idx, 0], 
           data[plot_idx, 1], 
           c=colors[plot_idx],
           alpha=0.4)

h = lambda c: plt.Line2D([],[],color=c, ls="",marker="o")
plt.legend(handles=[h(sc.cmap(sc.norm(i))) for i in range(len(labels))],
           labels=list(labels))
plt.show()

或者，您可以使用特殊的分散处理程序，如问题 Why doesn't the color of the points in a scatter plot match the color of the points in the corresponding legend? 中所示。但这在这里似乎有点过分了。