Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python Scatter Plot with Colorbar and Legend Issues

I'm working with a pretty simple example. I create three scatter plots on the same set of axes, and each data set I plot has a different associated colormap. However, the legend does not look as I'd want it to; why is this?

import numpy as np
import matplotlib.pyplot as plt
t = np.linspace(0,10,100)
x = np.random.rand(100,3)
y = np.random.rand(100,3)

colmaps = ['Blues', 'Greys', 'Reds']
for i in range(3):
    plt.scatter(x[:,i], y[:,i], c=t, cmap=colmaps[i], label=i)

plt.legend()
plt.show()

This produces a figure like below: Legend entries are all blue

I was hoping for the first label to be blue, the second to be grey, and the third to be red, so they're associated with the colormap, but it looks like that's not how it works. Is there a simple way to do this?

Thanks

like image 655
Chimi Avatar asked Mar 16 '23 19:03

Chimi


1 Answers

You can set the legend colors as such:

import numpy as np
import matplotlib.pyplot as plt

t = np.linspace(0,10,100)
x = np.random.rand(100,3)
y = np.random.rand(100,3)
colmaps = ['Blues', 'Greys', 'Reds']

for i in range(3):
    plt.scatter(x[:,i], y[:,i], c=t, cmap=colmaps[i], label=i)

plt.legend()
ax = plt.gca()
legend = ax.get_legend()
legend.legendHandles[0].set_color(plt.cm.Blues(.8))
legend.legendHandles[1].set_color(plt.cm.Greys(.8))
legend.legendHandles[2].set_color(plt.cm.Reds(.8))
plt.show()

I set the color of each legendHandle to a specific value in the respective colormap.

enter image description here

If you make the scatter dot's size larger you can see the color and associate individual dots with the legend easier. I also set just one dot per scatter plot in the legend, rather than the default 3, and set the legend's alpha to 0.5, and alpha of scatter plot to 0.7.

...
for i in range(3):
    plt.scatter(x[:,i], y[:,i], c=t, cmap=colmaps[i], label=i, s=200, alpha=0.7)
plt.legend(markerscale=0.7, scatterpoints=1)
ax = plt.gca()
legend = ax.get_legend()
legend.legendHandles[0].set_color(plt.cm.Blues(.8))
legend.legendHandles[1].set_color(plt.cm.Greys(.8))
legend.legendHandles[2].set_color(plt.cm.Reds(.8))
legend.get_frame().set_alpha(0.5)
...

enter image description here

like image 160
Scott Avatar answered Mar 19 '23 06:03

Scott