CODEWITHSUNDEEP
pythonmatplotlib
kilojoules
kilojoules

Reputation: 10073

pyplot legend for scatter plot colored by values

When I use a variable for coloring a scatter plot, how can I make a legend stating what colors represent? How can I make the legend show a label of 0 represents empty and 1 represents full? 

import matplotlib.pyplot as plt
X = [1,2,3,1,2,3,4]
Y = [1,1,1,2,2,2,2]
label = [0,1,1,0,0,1,1]
plt.scatter(X, Y, c= label, s=50)
plt.show()

Upvotes: 1

Views: 1742

Answers (1)

Tonechas
Tonechas

Reputation: 13723

Give this code a try:

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
X = [1, 2 ,3, 1, 2, 3, 4]
Y = [1, 1, 1, 2, 2, 2, 2]
labels = [0, 1, 1, 0, 0, 1, 1]
key = {0: ('red', 'empty'), 1: ('green', 'full')}
plt.scatter(X, Y, c=[key[index][0] for index in labels], s=50)
patches = [mpatches.Patch(color=color, label=label) for color, label in key.values()]
plt.legend(handles=patches, labels=[label for _, label in key.values()], bbox_to_anchor=(1, .3))
plt.show()

And this is what you'll get:

enter image description here

To use colors or labels different than those shown in the figure you simply need to change the values of the dictionary key appropriately.

Upvotes: 2

Related Questions