A simple matplotlib plotting script

simple-plotting.py

#!/usr/bin/env python3

"""
A simple matplotlib plotting script

Use the `--non-interactive` option to generate figures on remote servers without
a graphical user interface (GUI). For more information on matplotlib's backends
see: https://matplotlib.org/faq/usage_faq.html#what-is-a-backend
"""

import os
import argparse
import numpy as np

if __name__ == '__main__':
    # Parse command-line arguments
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument(      '--non-interactive', dest='interactive',
                                                   default=True,
                                                   action='store_false')
    parser.add_argument('-s', '--save-fig', default=False, action='store_true')
    parser.add_argument('-t', '--fig-type', default='pdf', type=str)
    args = parser.parse_args()

    # Load matplotlib: If run in non-interactive mode, set the matplotlib
    # backend to 'agg' for rendering without a GUI (this needs to be done before
    # pyplot is imported)
    import matplotlib
    if (not args.interactive):
        matplotlib.use('agg', warn=True)
    import matplotlib.pyplot as plt
    print('matplotlib backend: %s' % (matplotlib.get_backend()))

    # Close leftover figures if in interactive mode
    if args.interactive:
        plt.close('all')

    # Always save the figure if in non-interactive mode
    if (not args.interactive):
        args.save_fig = True

    # Generate some data
    t = np.linspace(0, 2*np.pi, 101)
    a = np.sin(t)

    # Create the figure and axis objects
    fig, ax = plt.subplots()

    # Plot the data
    ax.plot(t, a, 'k-')
    ax.set_xlabel(r'$t$')
    ax.set_ylabel(r'$A$')
    fig.tight_layout()

    # Save the figure to a file
    fig_dir = '.'
    fig_name = 'sine-function'
    if args.save_fig:
        fig_fname = os.path.join(fig_dir, fig_name + '.' + args.fig_type)
        print('Save figure to: %s' % (fig_fname))
        fig.savefig(fig_fname)

    # Show the figure if in interactive mode
    if args.interactive:
        fig.show()