Search code examples
pythonmatplotliblegendsubplotbox

python : using a subplot as legend for the others subplots


I have the script below that generates a plot consisted of 5 histograms (5 subplots) :

import matplotlib.pyplot as plt
import numpy as np

from scipy.stats import norm
from scipy.optimize import curve_fit

mean_o = list()
sigma_o = list()
y3 = list()

### generate some data
for i in range( 5 ):
    y3.append( norm.rvs( size=150000 ) )
y3 = np.transpose( y3 )

for i in range(5):
    mean_o.append( np.mean( y3[ :, i ] ) )
    sigma_o.append(  np.std( y3[ :, i ] ) )

## Histograms
# Number of bins
Nbins=100
binwidth = np.zeros(5)

# Fitting curves
def gaussian( x, a , mean, sigma ):
  return a * np.exp( -( ( x - mean )**2 / (2 * sigma**2 ) ) )

fig = plt.figure()
ax = { i : fig.add_subplot( 2, 3, i+1) for i in range( 5 ) }


for i in range(5):
    ymin = min(y3[:,i])
    ymax = max(y3[:,i])
    binwidth[i] = ( ymax - ymin) / Nbins
    bins_plot = np.arange( ymin, ymax + binwidth[i], binwidth[i])
    histdata = ax[i].hist(
        y3[:,i],
        bins=bins_plot,
        label='bin '+str(i)
    )

    range_fit = np.linspace( ymin, ymax, 250)
    # Fitting and plot version 1
    popt, pcov = curve_fit(
        gaussian,
        0.5 * ( histdata[1][:-1] + histdata[1][1:] ),
        histdata[0],
        p0=( max(histdata[0]), mean_o[i], sigma_o[i] ) )
    ax[i].plot(range_fit, gaussian( range_fit, *popt ) )
    ax[i].axvline( x=mean_o[i], ls=':', c='r' )

    # Fitting and plot version 2
    params = norm.fit( y3[ ::, i ], loc=mean_o[i], scale=sigma_o[i] )
    nth = gaussian(
        range_fit,
        len( y3[::, i]) * binwidth[i] / np.sqrt( 2 * np.pi ),
        *params
    )
    ax[i].plot(range_fit, nth, ls="--" )
plt.tight_layout()
plt.show()

that generates the following figure :

initial plot

Now, I would like to use in the 6th subplot on the right bottom as a legend for the 5 other subplots where the best fit values appear with a corresponding box and value to each subplot (like a color code) :

wanted plot : using subplot as legend

I don't know at all how to use a subplot as a legend, Does anyone get an idea ?


Solution

  • You can create an extra subplot and put a custom legend in it. The custom legend could use dummy handles and a list of labels. To customize how the legend looks, many parameters can be used. ax.axes('off') will switch off the unneeded surrounding axes of the last subplot.

    Here is some example code. To simplify the calculations, y3 is represented by a 2D numpy array.

    import matplotlib.pyplot as plt
    import numpy as np
    from scipy.stats import norm
    from scipy.optimize import curve_fit
    
    ### generate some data
    y3 = norm.rvs(size=(5, 5000))
    mean_o = np.mean(y3, axis=1)
    sigma_o = np.std(y3, axis=1)
    
    ## Histograms
    # Number of bins
    Nbins = 100
    
    # Fitting curves
    def gaussian(x, a, mean, sigma):
         return a * np.exp(-((x - mean) ** 2 / (2 * sigma ** 2)))
    
    
    fig = plt.figure(figsize=(15, 7))
    ax = {i: fig.add_subplot(2, 3, i + 1) for i in range(5)}
    
    ymin = y3.min(axis=1)
    ymax = y3.max(axis=1)
    legend_labels = []
    
    for i in range(5):
         bins_plot = np.linspace(ymin[i], ymax[i], Nbins + 1)
         values, bins, patches = ax[i].hist(
              y3[i],
              bins=bins_plot,
              label='bin ' + str(i))
    
         range_fit = np.linspace(ymin[i], ymax[i], 250)
         # Fitting and plot version 1
         popt, pcov = curve_fit(
              gaussian,
              0.5 * (bins[:-1] + bins[1:]),
              values,
              p0=(values.max(), mean_o[i], sigma_o[i]))
         ax[i].plot(range_fit, gaussian(range_fit, *popt))
         ax[i].axvline(x=mean_o[i], ls=':', c='r')
    
         # Fitting and plot version 2
         params = norm.fit(y3[::, i], loc=mean_o[i], scale=sigma_o[i])
         nth = gaussian(
              range_fit,
              len(y3[::, i]) * binwidth[i] / np.sqrt(2 * np.pi),
              *params
         )
         ax[i].plot(range_fit, nth, ls="--")
         legend_labels.append(f'bin {i + 1}: {mean_o[i]:.5f}')
    ax5 = fig.add_subplot(2, 3, 6)
    dummy_handle = plt.Rectangle((0, 0), 0, 0, color='none')
    ax5.legend(handles=[dummy_handle] * len(legend_labels), labels=legend_labels,
               handlelength=0, handletextpad=0, fontsize=24,
               loc='center', shadow=True, facecolor='lightgrey')
    ax5.axis('off')
    
    plt.tight_layout()
    plt.show()
    

    custom legend in subplot