David
David

Reputation: 433

Add a legend to a figure

Here is a code

import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
import numpy as np
 
fig, subs = plt.subplots(4,3) #setting the shape of the figure in one line as opposed to creating 12 variables 
 
iris = load_iris() ##code as per the example 
data = np.array(iris['data'])
targets = np.array(iris['target'])
 
cd = {0:'r',1:'b',2:"g"}
cols = np.array([cd[target] for target in targets])

 
# Row 1 
 
subs[0][0].scatter(data[:,0], data[:,1], c=cols)
subs[0][1].scatter(data[:,0], data[:,2], c=cols)
subs[0][2].scatter(data[:,0], data[:,3], c=cols)
 
# Row 2 
 
subs[1][0].scatter(data[:,1], data[:,0], c=cols)
subs[1][1].scatter(data[:,1], data[:,2], c=cols)
subs[1][2].scatter(data[:,1], data[:,3], c=cols)
 
# Row 3 
 
subs[2][0].scatter(data[:,2], data[:,0], c=cols)
subs[2][1].scatter(data[:,2], data[:,1], c=cols)
subs[2][2].scatter(data[:,2], data[:,3], c=cols)
 
#Row 4 
 
subs[3][0].scatter(data[:,3], data[:,0], c=cols)
subs[3][1].scatter(data[:,3], data[:,1], c=cols)
subs[3][2].scatter(data[:,3], data[:,2], c=cols)
 
plt.show()

Output

I would be interested in adding a legend indicating the red dots represent 'setosa', green dots 'versicolor' and blue dots 'virginica'. That legends would be at the bottom and center of the above picture. How can I do that?

I think I have to play with fig.legend, but I am not sure at all how to do that.

Upvotes: 0

Views: 555

Answers (2)

Léonard
Léonard

Reputation: 2630

You can loop over the targets in one of the subplots, and make the legend appear outside this plot. Here is what I obtained with your code:

multiple scatter plots legend

Here is the code:


import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
import numpy as np

fig, subs = plt.subplots(4,3, constrained_layout=True) #setting the shape of the figure in one line as opposed to creating 12 variables

iris = load_iris() ##code as per the example
data = np.array(iris['data'])
target_names = iris['target_names']
targets = np.array(iris['target'])

cd = {0:'r',1:'b',2:"g"}
cols = np.array([cd[target] for target in targets])


# Row 1

subs[0][0].scatter(data[:,0], data[:,1], c=cols)
subs[0][1].scatter(data[:,0], data[:,2], c=cols)
subs[0][2].scatter(data[:,0], data[:,3], c=cols)

# Row 2

subs[1][0].scatter(data[:,1], data[:,0], c=cols)
subs[1][1].scatter(data[:,1], data[:,2], c=cols)
subs[1][2].scatter(data[:,1], data[:,3], c=cols)

# Row 3

subs[2][0].scatter(data[:,2], data[:,0], c=cols)
subs[2][1].scatter(data[:,2], data[:,1], c=cols)
subs[2][2].scatter(data[:,2], data[:,3], c=cols)

# Row 4
subs[3][0].scatter(data[:,3], data[:,0], c=cols)

# loop for central subplot at last row
for t, name in zip(np.unique(targets), target_names):
    subs[3][1].scatter(data[targets==t,3], data[targets==t,1], c=cd[t], label=name)
subs[3][1].legend(bbox_to_anchor=(2, -.2), ncol=len(target_names))  # you can play with bbox_to_anchor for legend position

subs[3][2].scatter(data[:,3], data[:,2], c=cols)


plt.savefig('legend')

EDIT: I have also found this post in the matplotlib documentation, where you can extract the scatter elements from the scatter plot directly (without using a for loop). I have tried on the IRIS dataset without being able to make it work.

Upvotes: 1

mullinscr
mullinscr

Reputation: 1748

Add label='versicor' etc to just one of your subplots:

# Row 1 
 
subs[0][0].scatter(data[:,0], data[:,1], c=cols, label='virginica')
subs[0][1].scatter(data[:,0], data[:,2], c=cols, label='setosa')
subs[0][2].scatter(data[:,0], data[:,3], c=cols, label='versicolor')

Then you can call fig.legend(loc='lower center', ncol=3) before plt.show().

I have set ncol=3 to make it short and wide.

See my example below:

import numpy as np
import matplotlib.pyplot as plt

a = range(1,11)
b = np.random.randn(10).cumsum()
c = np.random.randn(10).cumsum()
d = np.random.randn(10).cumsum()

fig, (ax1, ax2) = plt.subplots(2, figsize=(9,5))

ax1.plot(a,b, label = 'one')
ax1.plot(a,c, label = 'two')
ax1.plot(a,d, label = 'two')

ax2.plot(a,c)

fig.legend(loc='lower center', ncol=3)

plt.show()

Upvotes: 0

Related Questions