[FIXED] Add squares in the legend in matplotlib

Issue

I have the following figure I got from this code :

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import colors
import matplotlib.gridspec as gridspec
import matplotlib.ticker as plticker

cmap = matplotlib.colors.LinearSegmentedColormap.from_list('', ['#41fdfe', '#fffe7a', '#021bf9', '#363737', '#840000' , '#ad8150' , '#e6daa6' , '#89fe05'])
fig = plt.figure(figsize=(8, 6)) 
gs = gridspec.GridSpec(1, 2, width_ratios=[7, 1] , hspace = 6.0 , wspace = 0.0)  
ax0 = plt.subplot(gs[0])
ax0.imshow(data_first, interpolation='none', origin='lower', cmap=cmap)
color = ['brown', 'grey', '#d1b26f' , 'brown' , '#ad8150', 'grey' , 'blue']
for y in range(5):
    for x in range(3):
            if difference[y , x] >= 0.6:
              size = 1200
            elif difference[y , x] >= 0.4:
                size = 1000
            elif difference[y,x] >= 0.1:
              size = 600
            elif difference[y , x] <= 0.1:
              size = 300
            ax0.scatter(x, y, s=size
                        , c = color[int(data_second[y,x])], marker='s' , edgecolor = 'k')
   
loc = plticker.MultipleLocator(base=1.0) # this locator puts ticks at regular intervals
ax0.xaxis.set_major_locator(loc)
ax0.yaxis.set_major_locator(loc)
cb_colors = ["#41fdfe", "blue", "brown" , '#d1b26f','#ad8150' , "grey"]
num_colors = len(cb_colors)
cmap_ = matplotlib.colors.ListedColormap(cb_colors)
cax = fig.add_axes([0.7, 0.3, 0.1, 0.5])
cb = matplotlib.colorbar.ColorbarBase(cax, orientation='vertical',
                           cmap=cmap_, norm=plt.Normalize( - 0.5 , num_colors - 0.5 ))
cb.set_ticks(range(num_colors))
cb.ax.set_yticklabels(["Cumulative rainfall", "Distribution of \n the rainfall event", "SM: layer 3 ","SM : layer 2 " , "SM : layer 1" , "Snow melt"])

How can add squares of different sizes at the bottom of the figure (as shown in the red rectangle)?

Many thanks

Solution

@zorals, you can modify this code and apply it to your own code.

import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import numpy as np

labels = ['< 0.1' , "< 0.4","> 0.6"]
msizes = np.array([20, 25,30,])
for ii in range(3):
    plt.plot([],[], 's', markeredgecolor='k',markerfacecolor='none',markersize=msizes[ii],label = labels[ii])

leg = plt.legend( ncol = 3,
          fancybox=False,title="Difference:\n",title_fontsize=15,edgecolor='k',framealpha=1,
                fontsize=12, loc='lower left', bbox_to_anchor=(1, 0.0), borderaxespad=0.,
                scatterpoints=1,borderpad=1.5)
leg._legend_box.align = "left"
plt.show()

Answered By - Li Yupeng