(Header photo by Raimond Klavins via Unsplash)

This post reviews techniques for creating multi-component figures using matplotlib, and provides an overview of a provisional (as of release 3.6.2) API for the matplotlib.pyplot.subplot_mosaic() method.

.   .   .

Mix-and-match subplots with 

GridSpec and subplot_mosaic()

matplotlib is a versatile python data-visualization package, which accounts for its extensive use in data science as well as its central place in (Python-based) data science education. At the time of writing, it is over ten years old, and the most recent release is 3.6.2. The original idea behind matplotlib was to emulate some of the functionality present in MATLAB, a proprietary software package sold by The Mathworks, Inc. [1]. 

matplotlib allows users to customize parameters ranging from the color of each bar in a bar chart to the format of the head and tail of annotative arrows. Unfortunately, this deep customizability and extensibility comes alongside a complex and sometimes unintuitive API. One common critique of matplotlib is its steep learning curve when moving beyond basic built-in plots.

The extensibility of matplotlib, combined with its oftentimes cumbersome API likely contributed to the creation of other Python visualization packages (e.g. seaborn, plotnine/ggplot2, and pandas) that wrap the heavy-lifting performed by matplotlib in a more user-friendly API [2].

In my projects, I prefer to use seaborn for generating charts, since I believe its API is more intuitive and it produces higher quality first-pass graphics. However, if I want to perform specific formatting changes or add annotations to my figures, I often need to “peel back the wrapper” and interact with the matplotlib code that seaborn is based on.

Examples of visual storytelling with multiple axes

Effective infographics often combine several different types of data presentation together into a single figure. This technique is used in respected publications such as The Economist, and FiveThirtyEight to great effect.

In the Economist article entitled, “A record of ancient Rome’s economy turns up in a glacier in Greenland”, the authors present a figure that connects the concentration of lead in Greenlandic ice cores to historical events in the Iberian peninsula (the source of most silver in ancient Europe) [3, 4].

FiveThirtyEight’s article entitled, “Support For Gun Control Will Likely Rise After Uvalde. But History Suggests It Will Fade.”, by Geoffrey Skelley, Nathaniel Rakich and Elena Mejia utilize a significantly simpler stacked line-graph visualization to clearly demonstrate the relationship between news coverage of school shootings and support for stricter gun-control laws (this graphic, along with many FiveThirtyEight graphics, also makes good use of color, which I will address in a future article) [5].

Creating multi-axis figures with matplotlib

In the body of this article, I’ll go over some of the existing ways to create multi-axis figures using matplotlib and seaborn before introducing a more flexible method using the recently added (as of release 3.3) subplot_mosaic() technique [6]. There are many different ways to create multiple-axis plots in Python, but some are simpler than others. Below, I cover the following techniques:

•  plt.subplot() 

•  plt.subplots() 

•  plt.GridSpec + plt.add_subplot() 

•  plt.subplot_mosaic() 

Let’s start by importing our packages and borrowing a helpful annotation function from matplotlib’s “GridSpec demo” [7].

import matplotlib.pyplot as plt

import seaborn as sns

import pandas as pd

import numpy as np

"""This annotation function is borrowed from matplotlib's 'GridSpec demo'

It numbers the axes within a figure to make them easier to identify"""

def annotate_axes(fig):

    for i, ax in enumerate(fig.axes):

        ax.text(0.5, 0.5, "ax%d" % (i+1), va="center", ha="center")

        ax.tick_params(labelbottom=False, labelleft=False)

The Original —  plt.subplot()

This technique likely needs no explanation, but just in case, this article is a nice primer. The following code produces figure 1.

The Slightly Automated — plt.subplots()

This technique returns an array of Axes objects that can be used to address each subplot individually. The following code reproduces figure 1.

N_ROWS, N_COLS = (3, 2)

fig, axes = \

    plt.subplots(

        nrows=N_ROWS,

        ncols=N_COLS)

annotate_axes(fig)

plt.tight_layout()

The calls in the two code blocks above create rectangular arrays of equally sized subplots. Sometimes, equally-sized grid-aligned plots can adequately present the connection between different plots. More often, we can create spatial relationships that further enhance the presentation of our data using more flexible calls (i.e. more arguments and more classes).

The Unwrapped Versatility — plt.GridSpec and plt.add_subplot()

plt.GridSpec is the behind-the-scenes class wrapped in more accessible methods like plt.subplot() and plt.subplots(). Functionally, plt.GridSpec provides a plt.Figure object with the instructions needed to arrange other objects (such as plt.Axes) within a figure.

By itself, a GridSpec object does not create any axes or subplots. It only defines spaces in a plot that can be used to arrange other objects — like the scaffolding you put up during construction. 

To start, here is the code needed to create a four by four GridSpec:

N_ROWS, N_COLS = (4, 4)

fig = plt.figure(

    figsize=(6, 6),

    layout="tight")

grid = plt.GridSpec(

    nrows=N_ROWS,

    ncols=N_COLS,

    figure=fig)

We can access the elements of a GridSpec object like we would a two-dimensional numpy array to select a single cell or multiple contiguous cells:

grid[0, 0]    # The cell in the first row and first column

grid[1, :]    # The entire second row

grid[0, 1:4]  # The second through fourth column of the first row

grid[2:4, :]  # All columns of the second and third row

n.b. — GridSpec indexing assigns [0, 0] to the upper-left corner and ranges are inclusive of the start position and exclusive of the end position ([a, b))

With the “scaffolding” in place, we can pass elements of the grid to the add_subplot() (a method that must be called from a parent plt.Figure object) to create new Axes objects in the specified spaces. Calling the add_subplot() method using the GridSpec selections above, we produce figure 2.

ax1 = fig.add_subplot(grid[0, 0])

ax2 = fig.add_subplot(grid[1, :])

ax3 = fig.add_subplot(grid[0, 1:4])

ax4 = fig.add_subplot(grid[2:4, :])

Each of these four axes can be addressed individually by matplotlib, seaborn, or another plotting package by passing the corresponding plt.Axes object to the plotting method.

The New(ish)comer — plt.subplot_mosaic()

In the matplotlib 3.3.0 (July 16th, 2020) release, the development team introduced a new wrapper for plt.GridSpec designed to accept either ASCII art or nested lists to quickly generate complex subplot grids [6]. If you need a reminder of what ASCII art is, please refer to the following code block (generated using Patorjk’s TAAG app) [8]:

           _           _       _                               _       ____  

          | |         | |     | |                             (_)     / /\ \ 

 ___ _   _| |__  _ __ | | ___ | |_   _ __ ___   ___  ___  __ _ _  ___| |  | |

/ __| | | | '_ \| '_ \| |/ _ \| __| | '_ ` _ \ / _ \/ __|/ _` | |/ __| |  | |

\__ \ |_| | |_) | |_) | | (_) | |_  | | | | | | (_) \__ \ (_| | | (__| |  | |

|___/\__,_|_.__/| .__/|_|\___/ \__| |_| |_| |_|\___/|___/\__,_|_|\___| |  | |

                | |             ______                                \_\/_/ 

                |_|            |______|

The good news is that you do not have to be an ASCII artist to make use of the plt.subplot_mosaic() API. All you need to do is imagine the layout you would like for your final figure, and choose a letter to represent each of the subplots you would like on the figure. For example, the following plt.subplot_mosaic() call would reproduce figure 2:

ascii = """

   .AAAA

   BCCCC

   BCCCC

   BCCCC

"""

kwargs = {"figsize": (6, 6)}

fig, ax = \

   plt.subplot_mosaic(

       mosaic=ascii,

       **kwargs)

annotate_axes(fig)

If you compare the mosaic argument above to how we defined the GridSpec for figure 2, you can see the same overall structure 

• A four-by-four grid of characters: (nrows, ncols) = (4, 4)

• Four different subplots: A, B, C, D

• One small subplot in the upper-left corner: A

• One wide subplot in the first row: BBB

• A full-width subplot in the second row: CCCC

• A large, two row subplot at the bottom: DDDD\nDDDD

In my experience, using this ASCII-art approach to figure layout allows me to make much faster adjustments to the positions and sizes of subplots than altering a GridSpec I created by hand. 

Here is an example of a different figure morphology, along with the full code used to create it with plt.subplot_mosaic():

ascii = """

    .AAAA

    BCCCC

    BCCCC

    BCCCC

"""

kwargs = {"figsize": (6, 6)}

fig, ax = \

    plt.subplot_mosaic(

        mosaic=ascii,

        **kwargs)

annotate_axes(fig)

Figure 3 is an example of a marginal plot that is sometimes used to help present two-variable histograms by plotting the two-dimensional histogram on a central axis and then projecting the one-dimensional histogram for one variable on the two other axes. 

n.b. You can use . in the mosaic argument to leave white space in the final figure.

For comparison’s sake, here is the complete code needed to generate figure 3 using GridSpec:

N_ROWS, N_COLS = (4, 4)

fig = plt.figure(

    figsize=(6, 6),

    layout="tight")

grid = plt.GridSpec(

    nrows=N_ROWS,

    ncols=N_COLS,

    figure=fig)

ax1 = fig.add_subplot(grid[0, 1:4])

ax2 = fig.add_subplot(grid[1:4, 0])

ax3 = fig.add_subplot(grid[1:4, 1:4])

annotate_axes(fig)

Every step is more streamlined in the mosaic technique. With GridSpec, we create a figure, apply that grid to the figure, and then initialize each subplot with a different add_subplot() call using (in my opinion) unintuitive figure indexing. With subplot_mosaic() we create a figure and a dictionary of axes in a single call, using an easy-to-adjust string.

An important difference between plt.subplot_mosaic() and other subplot generation methods is the return type. plt.subplot_mosaic() returns a tuple with type (plt.Figure, dict[plt.Axes]). The dictionary is indexed based on the letters you used in the mosaic argument, so the dictionary key "A" would correspond to ax1, “B” to ax2, and "C" to ax3 (I think this is a much easier way to handle multiple axes than having many individual axes variables floating around).

Recap

In this article, we briefly reviewed some of the techniques available within matplotlib to create and arrange subplots within figures. To review, we covered:

•  plt.subplot() 

•  plt.subplots() 

•  plt.GridSpec + plt.add_subplot() 

•  plt.subplot_mosaic() 

There are many more arguments and customizations that can be applied to further refine your multi-axis graphs, but those topics can wait for another day. For now, I hope this article has introduced you to new tools and new uses of existing tools to make your data visualizations more powerful.