Matplotlib Scatterplot Point Size Legend

Question

Below is a small amount of code that produces the attached image using matplotlib's scatterplot.

I'm trying to get a "legend" that that shows the size of several points and the corresponding "z-value".

Short of building it myself, is there anything like this? A "size" analogy to the colorbar?

import matplotlib.pyplot as plt
import numpy as np

fig = plt.figure(figsize=(8,6))
inset = fig.add_subplot(111)

np.random.seed(0) # so the image is reproducible
x1 = np.random.rand(30)
y1 = np.random.rand(30)
z1 = np.random.rand(30)


axis = inset.scatter(x1,y1,s=z1*100,c=z1,vmin=0,vmax=1)

inset.set_xlabel("X axis")
inset.set_ylabel("Y axis")

cbar = fig.colorbar(axis,ticks=[0,0.5,1])
cbar.ax.set_yticklabels(["Low","Medium","High"])

plt.savefig('scatterplot-zscale.png',bbox_inches='tight')

Answer 1

To you get a legend, you need to pass the keyword label when calling scatter on at least one data point. One way to do it is to select 3 representative points from your data, and add them to the plot again with labels.

import matplotlib.pyplot as plt
import numpy as np

np.random.seed(0) # so the image is reproducible
x1 = np.random.rand(30)
y1 = np.random.rand(30)
z1 = np.random.rand(30)

fig = plt.figure(figsize=(8,6))
inset = fig.add_subplot(111)
# i prefer no outlines on the dots, so edgecolors='none'
axis = inset.scatter(x1,y1,s=z1*100,c=z1,vmin=0,vmax=1,edgecolors='none')

inset.set_xlabel("X axis")
inset.set_ylabel("Y axis")

cbar = fig.colorbar(axis,ticks=[0,0.5,1])
cbar.ax.set_yticklabels(["Low","Medium","High"])

# here we step over the sorted data into 4 or 5 strides and select the 
# last 3 steps as a representative sample, this only works if your 
# data is fairly uniformly distributed
legend_sizes = np.sort(z1)[::len(z1)//4][-3:]

# get the indices for each of the legend sizes
indices = [np.where(z1==v)[0][0] for v in legend_sizes]

# plot each point again, and its value as a label
for i in indices:
    inset.scatter(x1[i],y1[i],s=100*z1[i],c=z1[i], vmin=0,vmax=1,edgecolors='none',
                  label='{:.2f}'.format(z1[i]))
# add the legend
inset.legend(scatterpoints=1)