pyplot.scatter
allows for passing to c=
an array that corresponds to groups, which will then color the points based on those groups. However, this seems to not support generating a legend without specifically plotting each group separately.
So, for example, a scatter plot with groups colored can be generated by iterating over the groups and plotting each separately:
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
feats = load_iris()['data']
target = load_iris()['target']
f, ax = plt.subplots(1)
for i in np.unique(target):
mask = target == i
plt.scatter(feats[mask, 0], feats[mask, 1], label=i)
ax.legend()
Which generates:
I can achieve a similar looking plot without iterating over each group though:
f, ax = plt.subplots(1)
ax.scatter(feats[:, 0], feats[:, 1], c=np.array(['C0', 'C1', 'C2'])[target])
But I cannot figure out a way to generate a corresponding legend with this second strategy. All of the examples I've come across iterate over the groups, which seems...less than ideal. I know I can manually generate a legend, but again that seems overly cumbersome.
See Question&Answers more detail:
os 与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…