letter_subplots.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# letter_subplots.py
# Jim Bagrow
# Last Modified: 2021-05-07

import numpy as np
import matplotlib.pyplot as plt


def letter_subplots(axes=None, letters=None, xoffset=-0.1, yoffset=1.0, **kwargs):
    """Add letters to the corners of subplots (panels). By default each axis is
    given an uppercase bold letter label placed in the upper-left corner.

    Args
        axes : list of pyplot ax objects. default plt.gcf().axes.

        letters : list of strings to use as labels, default ["A", "B", "C", ...]

        xoffset, yoffset : positions of each label relative to plot frame
          (default -0.1,1.0 = upper left margin). Can also be a list of
          offsets, in which case it should be the same length as the number of
          axes.

        Other keyword arguments will be passed to annotate() when panel letters
        are added.

    Returns:
        list of strings for each label added to the axes

    Examples:
        Defaults:
            >>> fig, axes = plt.subplots(1,3)
            >>> letter_subplots() # boldfaced A, B, C
        
        Common labeling schemes inferred from the first letter:
            >>> fig, axes = plt.subplots(1,4)        
            >>> letter_subplots(letters='(a)') # panels labeled (a), (b), (c), (d)

        Fully custom lettering:
            >>> fig, axes = plt.subplots(2,1)
            >>> letter_subplots(axes, letters=['(a.1)', '(b.2)'], fontweight='normal')

        Per-axis offsets:
            >>> fig, axes = plt.subplots(1,2)
            >>> letter_subplots(axes, xoffset=[-0.1, -0.15])
            
        Matrix of axes:
            >>> fig, axes = plt.subplots(2,2, sharex=True, sharey=True)
            >>> letter_subplots(fig.axes) # fig.axes is a list when axes is a 2x2 matrix
    """

    # get axes:
    if axes is None:
        axes = plt.gcf().axes
    # handle single axes:
    try:
        iter(axes)
    except TypeError:
        axes = [axes]

    # set up letter defaults (and corresponding fontweight):
    fontweight = "bold"
    ulets = list('ABCDEFGHIJKLMNOPQRSTUVWXYZ'[:len(axes)])
    llets = list('abcdefghijklmnopqrstuvwxyz'[:len(axes)])
    if letters is None or letters == "A":
        letters = ulets
    elif letters == "(a)":
        letters = [ "({})".format(lett) for lett in llets ]
        fontweight = "normal"
    elif letters == "(A)":
        letters = [ "({})".format(lett) for lett in ulets ]
        fontweight = "normal"
    elif letters in ("lower", "lowercase", "a"):
        letters = llets

    # make sure there are x and y offsets for each ax in axes:
    if isinstance(xoffset, (int, float)):
        xoffset = [xoffset]*len(axes)
    else:
        assert len(xoffset) == len(axes)
    if isinstance(yoffset, (int, float)):
        yoffset = [yoffset]*len(axes)
    else:
        assert len(yoffset) == len(axes)

    # defaults for annotate (kwargs is second so it can overwrite these defaults):
    my_defaults = dict(fontweight=fontweight, fontsize='large', ha="center",
                       va='center', xycoords='axes fraction', annotation_clip=False)
    kwargs = dict( list(my_defaults.items()) + list(kwargs.items()))

    list_txts = []
    for ax,lbl,xoff,yoff in zip(axes,letters,xoffset,yoffset):
        t = ax.annotate(lbl, xy=(xoff,yoff), **kwargs)
        list_txts.append(t)
    return list_txts


if __name__ == '__main__':
    x1 = np.random.randn(100,)
    y1 = x1 + 0.1*np.random.randn(100,)
    y2 = np.sin(x1) + 0.1*np.random.randn(100,)

    fig,axes = plt.subplots(1,2, figsize=(6.4*1.67,4.8))
    axes[0].plot(x1,y1, 'o')
    axes[1].plot(x1,y2, 'o')

    axes[0].set_xlabel("$x$")
    axes[1].set_xlabel("$x$")
    axes[0].set_ylabel("$y_1$")
    axes[1].set_ylabel("$y_2$")

    #letter_subplots() # bold upper-case, like Science uses
    #letter_subplots(letters="a"), # bold lowercase, like Nature uses
    letter_subplots(letters="(a)") # parenthetical letters, like many math and eng venues use

    plt.tight_layout()
    plt.show()