Python matplotlib: insets and aligned legends

Sometimes it is needed to create a plot with several lines and to assign a legend to each of them, which quite often ends in cumbersome and clunky legends. In my opinion, a more elegant way is include the legend as aligned text aside to each plot line, which is possible using Python and matplotlib.

Furthermore, to better highlight some details it is also possible to use an inset, hence to add a box which zooms a particular detail of the plot.

In the following is the Python code. The style is optimised for an IEEEtran journal.

import scipy as sc
import matplotlib.pyplot as plt
import matplotlib
from cycler import cycler
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
plt.close('all')

# Create a list of colours from the colormap viridis
color = plt.get_cmap('viridis')(sc.linspace(0, 1, 12))

# Set the settings to perform well in a IEEEtran document
params_IEEEtran = {'legend.fontsize': 10,
                   'axes.labelsize': 10,
                   'axes.titlesize': 10,
                   'xtick.labelsize': 9,
                   'ytick.labelsize': 9,
                   'mathtext.fontset': 'stix',
                   'font.family': 'Times New Roman',
                   'mathtext.rm': 'serif',
                   'mathtext.bf': 'serif:bold',
                   'mathtext.it': 'serif:italic',
                   'mathtext.sf': 'sans\\-serif',
                   'grid.color': 'k',
                   'grid.linestyle': ':',
                   'grid.linewidth': 0.5,
                   'axes.xmargin': 0,
                   'axes.axisbelow': False,
                   'lines.linewidth': 1.0,
                   'legend.frameon': False,
                   'axes.prop_cycle': cycler('color', color),
                   'figure.figsize': [3, 2.4],
                   }
plt.rcParams.update(params_IEEEtran)

# Create some sample data
x = sc.linspace(0.1, 100)[:, sc.newaxis]
e = sc.arange(0, 11, 1)[sc.newaxis, :]
y = e*x

# Main plot
fig = plt.figure('FIG')
ax = fig.add_subplot(1, 1, 1)
ax.plot(x, y)

# Put labels outside the plot
xmax = max(x)
text = [ax.text(1.01*xmax, y[-1, x], r'$' + str(x) + r'x$',
                color=color[n]) for (n, x) in enumerate(e[0, :])]

# Create the inset
inset = fig.add_axes([0.25, 0.6, 0.3, 0.3])  # Fraction of figure size (3, 2.4)
inset.axes.get_xaxis().set_visible(False)
inset.axes.get_yaxis().set_visible(False)
inset.plot(x, y)
[x.set_linewidth(0.5) for x in inset.spines.values()]
inset.set_xlim(0, 5)
inset.set_ylim(-1, 50)

# Create the box for the inset
box, c1, c2 = mark_inset(ax, inset, loc1=4, loc2=3, lw=0.3,
                         fc="none", ec=matplotlib.colors.hex2color('#900000'),
                         zorder=200)
box.set_linewidth(1)
box.set_color('k')

# Set the axis labels
ax.set_xlabel(r'$x$')
ax.set_ylabel(r'$f(x)$')

fig.tight_layout()

# Save as png and pdf
fig.savefig('./example.pdf', bbox_inches='tight', pad_inches=0)
fig.savefig('./example.png', bbox_inches='tight', dpi=300, pad_inches=0)

Which will result in the following plot: