Create a tablet for the story

I made the following graph with the following code, and the data is here :

import numpy as np
import pylab as plt
from matplotlib import rc,rcParams
rc('text',usetex=True)
rcParams.update({'font.size':10})
import matplotlib.cm as cm
from matplotlib.ticker import NullFormatter
def plot(Z_s,CWL,filter_id,spectral_type,model_mag,mag,plot_name):
      f= [r"$U_{38}$",r"$B$"  ,r"$V$"  ,r"$R$"  , r"$I$" ,r"$MB416$",r"$MB461$",r"$MB485$",r"$MB518$",r"$MB571$",r"$MB604$",r"$MB646$",r"$MB696$",r"$MB753$",r"$MB815$",r"$MB856$",r"$MB914$"]
      wavetable=CWL/(1+Z_s)
      data=model_mag-mag 
      nplist=['E', 'Sbc', 'Scd', 'Irr', 'SB3', 'SB2']
      colors = cm.rainbow(np.linspace(0, 1, len(f)))
      FILTER=filter_id
      SED=spectral_type
      for (j,d) in enumerate(nplist):
          bf=(SED==j)   
          if (j<3):
             k=j
             i_subplot = k + 1
             fig = plt.figure(1, figsize=(6,6))
             ax = fig.add_subplot(3,1,i_subplot)
             for i in range(len(f)):
                 bb=np.where(FILTER[bf]==i)[0]
                 ax.scatter(wavetable[bb], data[bb],  s=1, color=colors[i],label=f[i])
             if (k<2):
                ax.xaxis.set_major_formatter( NullFormatter() )
                ax.set_ylabel(r'$\Delta$ MAG',fontsize=10)
             else:
                ax.set_xlabel(r'WL($\AA$)',fontsize=10)
                ax.set_ylabel(r'$\Delta$ MAG',fontsize=10)
                leg = ax.legend(loc='lower center',prop={'size':4}, ncol=5)
                leg.get_frame().set_edgecolor('white')
             fig.subplots_adjust(wspace=0,hspace=0)
             ax.axhline(y=0,color='k')
             ax.set_xlim(1000,9000)
             ax.set_ylim(-3,3)
             ax.set_xticks(np.linspace(1000, 9000, 16, endpoint=False))
             ax.set_yticks(np.linspace(-3, 3, 4, endpoint=False)) 
             ax.text(8500,2.1,nplist[j], {'color': 'k', 'fontsize': 10})
             fontsize=8
             for tick in ax.xaxis.get_major_ticks():
                 tick.label1.set_fontsize(fontsize)
             for tick in ax.yaxis.get_major_ticks():
                 tick.label1.set_fontsize(fontsize)
             if (j==2):
                fname = plot_name+'.'+nplist[0]+'.'+nplist[1]+'.'+nplist[2]+'.pdf'        
                plt.savefig(fname)
                plt.close()
          else:
             k=j-3
             i_subplot = k + 1
             fig = plt.figure(1, figsize=(6,6))
             ax = fig.add_subplot(3,1,i_subplot)
             for i in range(len(f)):
                 bb=np.where(FILTER[bf]==i)[0]
                 ax.scatter(wavetable[bb], data[bb],  s=1, color=colors[i],label=f[i])
             if (k<2):
                ax.xaxis.set_major_formatter( NullFormatter() )
                ax.set_ylabel(r'$\Delta$ MAG',fontsize=10)
             else:
                ax.set_xlabel(r'WL($\AA$)',fontsize=10)
                ax.set_ylabel(r'$\Delta$ MAG',fontsize=10)
                leg = ax.legend( loc='lower center',prop={'size':4}, ncol=5)
                leg.get_frame().set_edgecolor('white')
             fig.subplots_adjust(wspace=0,hspace=0)
             ax.axhline(y=0,color='k')
             ax.set_xlim(1000,9000)
             ax.set_ylim(-3,3)
             ax.set_xticks(np.linspace(1000, 9000, 16, endpoint=False))
             ax.set_yticks(np.linspace(-3, 3, 4, endpoint=False)) 
             ax.text(8500,2.1,nplist[j], {'color': 'k', 'fontsize': 10})
             fontsize=8
             for tick in ax.xaxis.get_major_ticks():
                 tick.label1.set_fontsize(fontsize)
             for tick in ax.yaxis.get_major_ticks():
                 tick.label1.set_fontsize(fontsize)
             if (j==5):
                fname = plot_name+'.'+nplist[3]+'.'+nplist[4]+'.'+nplist[5]+'.pdf'        
                plt.savefig(fname)
                plt.close()

a=np.loadtxt('calibration.photometry.information.capak.cat')
Z_s=a[:,0]
CWL=a[:,1]
filter_id=a[:,2]
spectral_type=a[:,3]
model_mag=a[:,4]
mag=a[:,5]
plot_name='test'
plot(Z_s,CWL,filter_id,spectral_type,model_mag,mag,plot_name)

      

        
        
        
      

    

Is it possible to change different color legends from what I already did (in the last subheading) to a line close to the image with the same size as the whole plot? I'm looking for something similar to a colormap for projected 3D plots in 2D that just displays a hard-coded number of legend values ​​in my case?