Labelled lines – prettier charts in matplotlib

The default legend for matplotlib line charts can leave a little to be desired. With many colours it can also sometimes be a little tricky to match the legend to the appropriate line. Suppose instead we place the labels next to the lines.

Smoothed Annual Population Change - WHO Population Data
Smoothed Annual Population Change – WHO Population Data

Additionally here, we’ve removed the top and right axes, increased the font sizes of the labels and set the ticks to extend outwards. We firstly take our pandas dataframe death_rates_data_frame and plot it as normal (disabling gridlines and making the lines slightly thicker). We add the larger axis labels.

population_data_frame.plot(legend=False, grid=False, linewidth=2, figsize=(9,6))
plt.xlabel('Year', fontsize=20)
plt.ylabel('Population \n% Change', rotation='horizontal', labelpad=80, fontsize=20)

Next we define a function that, given a set of axes, will

  • Hide the top and right axes
  • Set the ticks to only display on the bottom x-axis and the left y-axis
  • Set the ticks to extend outwards
  • Increase the font-size of the labels
def format_graph(ax):
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.get_xaxis().tick_bottom()
    ax.get_yaxis().tick_left()
    ax.tick_params(axis='both', direction='out', labelsize=14)

The label_lines function handles the placing of text next to the final vertex of the line.

def label_lines(ax, offset, label_formatter=None, **kwargs):
    for handle, label in zip(*ax.get_legend_handles_labels()):
        path = handle.get_path()
        #careful with the NaNs
        last_vertex = pd.DataFrame(path.vertices).dropna().as_matrix()[-1]
        nicer_label = label_formatter(label) if label_formatter else label
        plt.text(last_vertex[0]+offset[0], last_vertex[1]+offset[1], nicer_label, color=handle.get_color(), transform=ax.transData, **kwargs)

Finally we call the above code with the current axes! Given the slight label overlap, let’s outline the text in white with path_effects.

import matplotlib.patheffects as PathEffects
ax = plt.gca()
format_graph(ax)
label_lines(ax, offset=(1,0),
                fontsize=16,
                path_effects=[PathEffects.withStroke(linewidth=3, foreground="w")])

Voilà.

Advertisements
Labelled lines – prettier charts in matplotlib

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s