Functional Figure Composition in Matplotlib

TL;DR: This page describes an approach to composing figures in Python’s matplotlib. The writing is messy, but I think the main idea is clear enough to be useful.

Composing Figures in matplotlib can be challenging.

Options include

• subplots (Axes within a Figure, created from methods on Figure instances), and
• gridspecs (also subplots within a figure, created using a more expressive gridspec syntax).

These are ok, but there are some aspects of these solutions that I don’t particularly like.

Research Figures often need:

• Complex layouts - not always needed but nice to know it’s possible.
• Modularity - figures standalone (ie for a presentation slide) and composed (ie for a manuscript).
  • Often we want to maintain standalone Figure A, standalone Figure B, and a composed Figure A+B.
• Provenance & editability - the method described here doesn’t solve for this, but lays a foundation.

Existing Approaches for Composing Figures:

• If we go for composition of Axes within a Figure, I have never understood how to untangle one Axis from the rest of the Figure. It would be great to be able to do this so we a can access Figure API stuff, like savefig, to do things to just Figure A.
• The other choice is standalone Figures. It wasn’t clear to me how to compose these. The naive solution would be to collect a bunch of Figure objects then somehow stitch them together. Something like

figure_a = make_figure_a(data)
figure_b = make_figure_b(data)
 
plt.combine_figs([figure_a, figure_b])

I haven’t found a straightforward way to do this. Others have asked:

https://stackoverflow.com/questions/45810557/copy-an-axes-content-and-show-it-in-a-new-figure

Solutions in the stackoverflow post use pickle

...
inx = list(fig.axes).index(event.inaxes)
buf = io.BytesIO()
pickle.dump(fig, buf)
buf.seek(0)
fig2 = pickle.load(buf) 

for i, ax in enumerate(fig2.axes):
    if i != inx:
        fig2.delaxes(ax)
    else:
        axes=ax

axes.change_geometry(1,1,1)
fig2.show()
...

, or pull data out of the Figure object:

fig, axs = plot_something()
ax = axs[2]
l = list(ax.get_lines())[0]
l2 = list(ax.get_lines())[1]
l3 = list(ax.get_lines())[2]
plot(l.get_data()[0], l.get_data()[1])
plot(l2.get_data()[0], l2.get_data()[1])
plot(l3.get_data()[0], l3.get_data()[1])
ylim(0,1)

These solutions don’t feel very direct or readable.

• I don’t find the gridspec syntax very readable.

This is subjective, but here is an example of some gridspec code from the matplotlib docs:

import matplotlib.pyplot as plt

import matplotlib.gridspec as gridspec


def format_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)


# gridspec inside gridspec
fig = plt.figure()

gs0 = gridspec.GridSpec(1, 2, figure=fig)

gs00 = gridspec.GridSpecFromSubplotSpec(3, 3, subplot_spec=gs0[0])

ax1 = fig.add_subplot(gs00[:-1, :])
ax2 = fig.add_subplot(gs00[-1, :-1])
ax3 = fig.add_subplot(gs00[-1, -1])

# the following syntax does the same as the GridSpecFromSubplotSpec call above:
gs01 = gs0[1].subgridspec(3, 3)

ax4 = fig.add_subplot(gs01[:, :-1])
ax5 = fig.add_subplot(gs01[:-1, -1])
ax6 = fig.add_subplot(gs01[-1, -1])

plt.suptitle("GridSpec Inside GridSpec")
format_axes(fig)

plt.show()

A Different Approach

I want to highlight a less-emphasized way of composing figures in matplotlib. This isn’t new, but I don’t see it used or discussed often.

To do this, we’ll use some techniques from functional programming.

One of them is partial function application. Another is first-class functions.

This is also a chance to use some less-known features of Python typing, like Protocol.

Part 1: Background

Let’s read the docs here:

https://matplotlib.org/stable/gallery/subplots_axes_and_figures/subfigures.html

We can see there is a way to nest visualizations at the Figure level, instead of the Axes level.

https://matplotlib.org/stable/_images/sphx_glr_subfigures_004_2_00x.png

Part 2: Strategy

Instead of trying to compose figure objects, we’ll compose the functions that create them.

We’ll define functions that accept a Figure (expected to be empty), some data, and other keyword arguments, and return a Figure.

How to compose them: We need a function to orchestrate the composition of the other functions.

Let’s call this compose_on_subfigures. This is where first-class functions come in — we’ll pass this function a list of the visualizer functions from before.

compose_on_subfigures looks like this:

It takes a Figure, some data, some keyword arguments that dictate how to split up fig into SubFigures, and a list of subfigure_fns.

Each of the subfigure_fns is a Callable that should accept two arguments, fig and data, plus arbitrary keyword arguments, and return a Figure.

def compose_on_subfigures(
    fig: Figure,
    data: pd.DataFrame,
    subfigure_kwargs: dict,
    subfigure_fns: List[Callable],
) -> Figure:
    """
    This function creates a figure with multiple subfigures.
 
    Parameters:
    data: The data to be plotted.
    subfigure_kwargs (dict): The keyword arguments to be passed to the subfigures.
    subfigure_fns (List[Callable]): A list of functions that generate each subfigure.
 
    Returns:
    Figure: The final figure with all subfigures.
    """
    subfigs = fig.subfigures(**subfigure_kwargs)
    for subfig, subfig_fn in zip(subfigs, subfigure_fns):  # type: ignore
        subfig_fn(data=data, fig=subfig)
 
    return fig

Notice anything? The function signature for compose_on_subfigures, and the function signature for a subfig_fn are almost the same! This allows us to nest calls to compose_on_subfigures to achieve any nested grid layout we want. And because each of the subfigure_fns return Figures (instead of Axes), we can also use this completely standalone.

More precise typing with Protocol

If we want to enhance the level of specificity of typing for the subfigure_fns argument, we have some options:

Callable

or

Callable[Figure,DataFrame] -> Figure

The latter is pretty good. Some downsides though:

• it indicates typing by position. Not necessarily a downside but not clear whether it will support calls like my_figure_fn(data,fig) and my_figure_fn(fig,data) equally well. Haven’t checked but it probably shouldn’t.

Also, I don’t think **kwargs is supported in the signature for Callable.

For more strict typing, we can use Protocol from the typing module.

class FigureMaker(Protocol):
    def __call__(
        self,
        fig: Figure,
        data: pd.DataFrame
        **kwargs,
    ) -> Figure: ...

This type annotation can be used like:

def compose_on_subfigures(
    fig: Figure,
    data: pd.DataFrame,
    subfigure_kwargs: dict,
    subfigure_fns: List[FigureMaker],
) -> Figure:
    # ... same as above
    return fig