在Python中将图例添加到散点图

Question

I have made a scatter plot in the following way:

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 holds the data and y holds the labels (1,2,3). Now I would like to add a legend.

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

This does not work, it only prints label1. How can this be done otherwise?

Answer 1

The idea is to divide data set on separate data sets which are represented by the same color. After that the legend displayed properly.

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()