Python ProgrammingPython Programming

Create legend with bubble size using Numpy and Matplotlib

Legend with bubble size

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

N = 50
M = 5 # Number of bins

x = np.random.rand(N)
y = np.random.rand(N)
a2 = 400*np.random.rand(N)

# Create the DataFrame from your randomised data and bin it using groupby.
df = pd.DataFrame(data=dict(x=x, y=y, a2=a2))
bins = np.linspace(df.a2.min(), df.a2.max(), M)
grouped = df.groupby(np.digitize(df.a2, bins))

# Create some sizes and some labels.
sizes = [50*(i+1.) for i in range(M)]
labels = ['Tiny', 'Small', 'Medium', 'Large', 'Huge']

for i, (name, group) in enumerate(grouped):
    plt.scatter(group.x, group.y, s=sizes[i], alpha=0.5, label=labels[i])

plt.legend()
plt.show()

The following is the output that will be obtained:


Create legend with bubble size using Numpy and Matplotlib