Making a Matplotlib legend for a 3D Two-Axis Scatter Plot

0
votes

I'm making a scatterplot with matplotlib based on a dictionary, where the dictionary maps a two-element tuple to an integer value. The coordinate represented by the tuple is plotted, and then is colored based on the value of the integer it maps to. 

xs = [tup[0] for tup in list(mydict.keys())]
ys = [tup[1] for tup in list(mydict.keys())]
colors = ['g' if outcome > 0 else 'r' for outcome in list(mydict.values())]
plt.figure()
plt.subplot(211)
plt.scatter(xs, ys, c = colors, alpha = 0.5, s = 4)

I want to put a legend in the upper right corner of this scatter plot showing that a red dot means "Ineffective" and a green dot means "Effective." Since I draw both colored points in one line, I can't (or don't know how to) give the colors separate label parameters. So, how can I create the legend?

1
pythonmatplotlib

1 Answers

1
votes

Here's a simple solution.

xs = np.array([tup[0] for tup in list(mydict.keys())])
ys = np.array([tup[1] for tup in list(mydict.keys())])
outcomes = np.array([outcome for outcome in list(mydict.values())])
plt.figure()
plt.subplot(211)
plt.scatter(xs[outcomes==0], ys[outcomes==0], c = colors, alpha = 0.5, s = 4, label = 'Ineffective')
plt.scatter(xs[outcomes==1], ys[outcomes==1], c = colors, alpha = 0.5, s = 4, label = 'Effective')
plt.legend(loc='upper right')
plt.show()