Python plotting script

make-plots

#!/usr/bin/env python3

import argparse

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as tick


aspect = 1/1.618
resolution = 72.27
columnwidth = 246/resolution
textwidth = 510/resolution
textiny, texsmall, texnormal = 8.0, 9.25, 10.0
offblack = '#262626'

plt.rcdefaults()
plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['CMU Serif'],
    'font.size': texsmall,
    'legend.fontsize': texsmall,
    'axes.labelsize': texsmall,
    'axes.titlesize': texsmall,
    'xtick.labelsize': textiny,
    'ytick.labelsize': textiny,
    'font.weight': 400,
    'axes.labelweight': 400,
    'axes.titleweight': 400,
    'lines.linewidth': .9,
    'lines.markersize': 3,
    'lines.markeredgewidth': .1,
    'patch.linewidth': .9,
    'axes.linewidth': .5,
    'xtick.major.width': .5,
    'ytick.major.width': .5,
    'xtick.minor.width': .5,
    'ytick.minor.width': .5,
    'xtick.major.size': 2,
    'ytick.major.size': 2,
    'xtick.minor.size': 1.3,
    'ytick.minor.size': 1.3,
    'xtick.major.pad': 1.8,
    'ytick.major.pad': 1.8,
    'text.color': offblack,
    'axes.edgecolor': offblack,
    'axes.labelcolor': offblack,
    'xtick.color': offblack,
    'ytick.color': offblack,
    'legend.numpoints': 1,
    'legend.scatterpoints': 1,
    'legend.frameon': False,
    'image.interpolation': 'none',
    'pdf.fonttype': 42
})


plot_functions = {}


def plot(f):
    def wrapper(*args, **kwargs):
        print(f.__name__)
        f(*args, **kwargs)
        plt.savefig('{}.pdf'.format(f.__name__))
        plt.close()

    plot_functions[f.__name__] = wrapper

    return wrapper


def finish(despine=True, remove_ticks=False, pad=0.1, h_pad=None, w_pad=None):
    fig = plt.gcf()

    for ax in fig.axes:
        if despine:
            for spine in 'top', 'right':
                ax.spines[spine].set_visible(False)

        if remove_ticks:
            for ax_name in 'xaxis', 'yaxis':
                getattr(ax, ax_name).set_ticks_position('none')
        else:
            ax.xaxis.set_ticks_position('bottom')
            ax.yaxis.set_ticks_position('left')

    fig.tight_layout(pad=pad, h_pad=h_pad, w_pad=w_pad)


def set_loc(ax, xy=None, nbins=5, steps=[1, 2, 3, 4, 10],
            prune=None, minor=0):
    if xy == 'x':
        axes = ax.xaxis,
    elif xy == 'y':
        axes = ax.yaxis,
    else:
        axes = ax.xaxis, ax.yaxis

    for axis in axes:
        axis.set_major_locator(
            tick.MaxNLocator(nbins=nbins, steps=steps, prune=prune)
        )
        if minor:
            axis.set_minor_locator(tick.AutoMinorLocator(minor))


@plot
def test():
    plt.figure(figsize=(columnwidth, aspect*columnwidth))

    X = np.arange(10)
    plt.plot(X, X*X, '-o', color=plt.cm.Blues(0.8))

    finish()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('plots', nargs='*')
    args = parser.parse_args()

    if args.plots:
        for i in args.plots:
            if i.endswith('.pdf'):
                i = i[:-4]
            if i in plot_functions:
                plot_functions[i]()
            else:
                print('unknown plot:', i)
    else:
        for f in plot_functions.values():
            f()


if __name__ == "__main__":
    main()