CODEWITHSUNDEEP
pythonmatplotliblegend
machinery
machinery

Reputation: 6290

Adding legend to scatter plot in Python

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?

Upvotes: 0

Views: 1059

Answers (1)

Roman Fursenko
Roman Fursenko

Reputation: 688

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

enter image description here

Upvotes: 1

Related Questions