Search code examples
matplotliblegendcolormaplinestyle

How to generate several legends for single plot matplotlib

I was making a plot of f(x,y,z) and wanted this to be displayed in a 2D-plane. To avoid cluttering my legend i decided to have different linestyles for y, different colors for z and place the two in two separate legends. I couldn't find out how to do this even after a lot of digging, so I'm posting the solution i came up with here :) If anyone has more elegant solutions I'm all ears :)

Solution

  • Basically the solution was to make three plots, set two of them to have size (0,0) and place those two where i wanted the legends. It feels like an ugly way to do it, but it gave a nice plot and i didn't find any other way :) The resulting plot looks like this: Pretty plot

    def plot_alt(style = 'log'):

cmap = cm.get_cmap('inferno')
color_scale = 1.2 #Variable to get colors from a certain part of the colormap

#Making grids for delta T and average concentration
D_T_axis = -np.logspace(np.log10(400), np.log10(1), 7)
C_bar_list = np.linspace(5,10,4)

ST_list = np.logspace(-3,-1,100)

# f(x,y,z)
DC_func = lambda C_bar, ST, DT: 2*C_bar * (1 - np.exp(ST*DT))/(1 + np.exp(ST*DT))

#Some different linestyles
styles = ['-', '--', '-.', ':']

fig, ax = plt.subplots(1,3, figsize = (10,5))

plt.sca(ax[0])
for i, C_bar in enumerate(C_bar_list): #See plot_c_rel_av_DT() for 'enumerate'
    for j, DT in enumerate(D_T_axis):
        plt.plot(ST_list, DC_func(C_bar, ST_list, DT), color = cmap(np.log10(-DT)/(color_scale*np.log10(-D_T_axis[0]))),
                                   linestyle = styles[i])

# Generating separate legends by plotting lines in the two other subplots
# Basically: to get two separate legends i make two plots, place them where i want the legends
# and set their size to zero, then display their legends.
plt.sca(ax[1]) #Set current axes to ax[1]
for i, C_bar in enumerate(C_bar_list):
    # Plotting the different linestyles
    plt.plot(C_bar_list, linestyle = styles[i], color = 'black', label = str(round(C_bar, 2)))
plt.sca(ax[2])
for DT in D_T_axis:
    #plotting the different colors
    plt.plot(D_T_axis, color = cmap(np.log10(-DT)/(color_scale*np.log10(-D_T_axis[0]))), label = str(int(-DT)))

#Placing legend
#This is where i move and scale the three plots to make one plot and two legends
box0 = ax[0].get_position() #box0 is an object that contains the position and dimentions of the ax[0] subplot
box2 = ax[2].get_position()

ax[0].set_position([box0.x0, box0.y0, box2.x0 + 0.4*box2.width, box0.height])
box0 = ax[0].get_position()

ax[1].set_position([box0.x0 + box0.width, box0.y0 + box0.height + 0.015, 0,0])
ax[1].set_axis_off()
ax[2].set_position([box0.x0 + box0.width ,box0.y0 + box0.height - 0.25, 0,0])
ax[2].set_axis_off()

#Displaying plot
plt.sca(ax[0])

plt.xscale('log')
plt.xlim(0.001, 0.1)
plt.ylim(0, 5)

plt.xlabel(r'$S_T$')
plt.ylabel(r'$\Delta C$')
ax[1].legend(title = r'$\langle c \rangle$ [mol/L]',
             bbox_to_anchor = (1,1), loc = 'upper left')
ax[2].legend(title = r'$-\Delta T$ [K]', bbox_to_anchor = (1,1), loc = 'upper left')

#Suptitle is the title of the figure. You can also have titles for the individual subplots
plt.suptitle('Steady state concentration gradient as a function of Soret-coefficient\n'
             'for different temperature gradients and total concentrations')