Добавление легенды на точечную в Python

Question

Я сделал график рассеяния следующим образом:

f, ax1 = plt.subplots(3,2) 
cmap = matplotlib.cm.get_cmap('coolwarm') 
ax1[0,1].scatter(data[:,0], data[:,1], c=y, s=20, marker='o', alpha=.5, cmap=cmap) 

data содержит данные и y удерживает этикетки (1,2,3). Теперь я хотел бы добавить легенду.

ax1[0,1].legend(('label1', 'label2', 'label3'), 
      scatterpoints=1, 
      loc='lower left', 
      fontsize=10) 

Это не работает, оно печатает только этикетку1. Как это можно сделать иначе?

Answer 1

Идея состоит в том, чтобы разделить набор данных на отдельные наборы данных, которые представлены одним и тем же цветом. После этого легенда отображается правильно.

import matplotlib.pyplot as plt 
import matplotlib 
import numpy as np 

data = np.zeros(shape=(10,2)) 
data[:,0] = np.linspace(0,1,10) 
data[:,1] = np.linspace(0,1,10) 

y = ['red', 'green', 'blue'] 
f, ax1 = plt.subplots(3,2) 
cmap = matplotlib.cm.get_cmap('coolwarm') 
ny = len(y) 
for i, itm in enumerate(y): 
    datac = data[i::ny,:] 
    ax1[0,1].scatter(datac[:,0], datac[:,1], c=itm, 
        s=20, marker='o', alpha=.5, cmap=cmap) 

ax1[0,1].legend(['label1', 'label2', 'label3'], 
      scatterpoints=1, 
      loc='lower left', 
      fontsize=10) 
plt.show()