Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse files

Change to argparse module.

  • Loading branch information...
commit 1bec6c02b1b8f1a765357e4465b39b3427716566 1 parent b2e78d8
@lmjohns3 authored
Showing with 113 additions and 116 deletions.
  1. +113 −116 scripts/py-grep-plot
View
229 scripts/py-grep-plot
@@ -1,6 +1,6 @@
#!/usr/bin/env python
-# Copyright (c) 2011 Leif Johnson <leif@leifjohnson.net>
+# Copyright (c) 2011-2013 Leif Johnson <leif@leifjohnson.net>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
@@ -22,77 +22,80 @@
'''A command-line script for plotting data from text files.'''
+import argparse
import bz2
import glob
import gzip
import logging
import numpy
-import optparse
import os
import re
import sys
from matplotlib import pyplot
-FLAGS = optparse.OptionParser('Usage: py-grep-plot [OPTIONS] [FILE...]')
-FLAGS.add_option('-o', '--output', metavar='FILE',
- help='save to FILE instead of displaying on screen')
-
-g = optparse.OptionGroup(FLAGS, 'Data Options')
-g.add_option('-b', '--batch', type=int, metavar='N',
- help='batch data into groups of N points and plot mean + std')
-g.add_option('-e', '--every', type=int, metavar='N',
- help='restrict plot to show only every Nth data point')
-g.add_option('-k', '--column', action='append', type=int, metavar='K',
- help='extract data from the Kth space-separated column')
-g.add_option('-r', '--regex', default=r'([-+eE.\d]+)', metavar='RE',
- help='extract data points from FILEs with RE ([-+eE.\d]+)')
-g.add_option('-s', '--smooth', type=int, default=0, metavar='N',
- help='smooth across N points before plotting (0)')
-FLAGS.add_option_group(g)
-
-g = optparse.OptionGroup(FLAGS, 'Series Options')
-g.add_option('-a', '--alpha', type=float, default=0.9, metavar='N',
- help='plot series with alpha N (0.9)')
-g.add_option('-c', '--colors', default='k,r,#666666,c,b,m,g,y', metavar='S,S,...',
- help='cycle through the given colors (k,r,c,b,m,g,y)')
-g.add_option('-p', '--points', default='o-', metavar='S,S,...',
- help='cycle through the given line/point styles (o-)')
-FLAGS.add_option_group(g)
-
-g = optparse.OptionGroup(FLAGS, 'Plot Options')
-g.add_option('-g', '--grid', default=False, action='store_true',
- help='include a grid (False)')
-g.add_option('-L', '--legend', metavar='[tl|cl|bl|tr|cr|br]',
- help='include a legend (None)')
-g.add_option('-l', '--log', default='', metavar='[x|y|xy]',
- help='use a log scale on the specified axes')
-g.add_option('-t', '--title', metavar='S',
- help='use S as the plot title')
-g.add_option('-x', '--xlabel', metavar='S',
- help='use S as the label for the x-axis')
-g.add_option('-y', '--ylabel', metavar='S',
- help='use S as the label for the y-axis')
-g.add_option('-X', '--xlim', metavar='A,B',
- help='use (A,B) as the range for the x-axis')
-g.add_option('-Y', '--ylim', metavar='A,B',
- help='use (A,B) as the range for the y-axis')
-FLAGS.add_option_group(g)
-
-
LEGEND = {
- 'upper right': 1, 'ur': 1, 'tr': 1,
- 'upper left': 2, 'ul': 2, 'tl': 2,
- 'lower left': 3, 'll': 3, 'bl': 3,
- 'lower right': 4, 'lr': 4, 'br': 4,
- 'right': 5, 'r': 5,
- 'center left': 6, 'cl': 6, 'ml': 6,
- 'center right': 7, 'cr': 7, 'mr': 7,
- 'lower center': 8, 'lc': 8, 'bc': 8,
- 'upper center': 9, 'uc': 9, 'tc': 9,
- 'center': 10, 'c': 10,
+ 'ul': 2, 'tl': 2,
+ 'uc': 9, 'tc': 9,
+ 'ur': 1, 'tr': 1,
+ 'cl': 6,
+ 'cc': 10,
+ 'cr': 7,
+ 'll': 3, 'bl': 3,
+ 'lc': 8, 'bc': 8,
+ 'lr': 4, 'br': 4,
}
+FLAGS = argparse.ArgumentParser(
+ conflict_handler='resolve',
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+FLAGS.add_argument('input', metavar='FILE', nargs=argparse.REMAINDER)
+
+g = FLAGS.add_argument_group('output')
+g.add_argument('-o', '--output', type=argparse.FileType, metavar='FILE',
+ help='save to FILE instead of displaying on screen')
+g.add_argument('-D', '--dpi', type=int, metavar='N',
+ help='save figure with dpi N')
+
+g = FLAGS.add_argument_group('data')
+g.add_argument('-b', '--batch', type=int, metavar='N',
+ help='batch data into groups of N points and plot mean + std')
+g.add_argument('-e', '--every', type=int, metavar='N',
+ help='restrict plot to show only every Nth data point')
+g.add_argument('-k', '--column', action='append', type=int, metavar='K',
+ help='extract data from the Kth space-separated column')
+g.add_argument('-r', '--regex', default=r'([-+eE.\d]+)', metavar='RE',
+ help='extract data points from inputs using RE')
+g.add_argument('-s', '--smooth', type=int, metavar='N',
+ help='smooth across N points before plotting')
+
+g = FLAGS.add_argument_group('series')
+g.add_argument('-a', '--alpha', type=float, default=0.9, metavar='N',
+ help='plot series with alpha N')
+g.add_argument('-c', '--colors', nargs='+', default=tuple('krcbmgy'), metavar='C',
+ help='cycle through the given colors')
+g.add_argument('-p', '--points', nargs='+', default=['o-'], metavar='S',
+ help='cycle through the given line/point styles')
+
+g = FLAGS.add_argument_group('axes')
+g.add_argument('-g', '--grid', action='store_true',
+ help='include a grid')
+g.add_argument('-L', '--legend', choices=tuple(sorted(LEGEND.keys())),
+ help='include a legend (None)')
+g.add_argument('-l', '--log', choices=['x', 'y', 'xy'],
+ help='use a log scale on the specified axes')
+g.add_argument('-t', '--title', metavar='S',
+ help='use S as the plot title')
+g.add_argument('-x', '--xlabel', metavar='S',
+ help='use S as the label for the x-axis')
+g.add_argument('-y', '--ylabel', metavar='S',
+ help='use S as the label for the y-axis')
+g.add_argument('-X', '--xlim', metavar='A,B',
+ help='use (A,B) as the range for the x-axis')
+g.add_argument('-Y', '--ylim', metavar='A,B',
+ help='use (A,B) as the range for the y-axis')
+
def extract_columns(data, columns, x, y, ey):
'''Pull specific column values out to plot.'''
@@ -182,107 +185,101 @@ def read_input(args):
yield os.path.splitext(os.path.basename(filename))[0], handle
-def compile_regex(opts):
+def compile_regex(args):
'''Compile a regular expression pattern.'''
- logging.debug('compiling REGEX %r', opts.regex)
+ logging.debug('compiling REGEX %r', args.regex)
try:
- return re.compile(opts.regex)
+ return re.compile(args.regex)
except:
- logging.critical('cannot compile REGEX %r', opts.regex)
+ logging.critical('cannot compile REGEX %r', args.regex)
sys.exit(-2)
-def make_axes(opts):
+def make_axes(args):
'''Create an axes object to hold our plots.'''
- X, Y = opts.ylabel and 0.12 or 0.1, opts.xlabel and 0.13 or 0.1
+ X, Y = args.ylabel and 0.12 or 0.1, args.xlabel and 0.13 or 0.1
ax = pyplot.axes([X, Y, 0.95 - X, 0.95 - Y])
ax.xaxis.tick_bottom()
ax.yaxis.tick_left()
- if 'x' in opts.log:
+ if 'x' in args.log:
ax.set_xscale('log')
- if 'y' in opts.log:
+ if 'y' in args.log:
ax.set_yscale('log')
return ax
-def format_axes(ax, opts):
+def format_axes(ax, args):
'''Format our plotting axes using the script options.'''
- logging.debug('using legend: %s' % opts.legend)
- loc = LEGEND.get(opts.legend)
+ logging.debug('using legend: %s' % args.legend)
+ loc = LEGEND.get(args.legend)
if loc is not None:
ax.legend(loc=loc)
- logging.debug('using grid: %s' % opts.grid)
- ax.grid(opts.grid)
-
- if opts.title:
- logging.debug('using title: %r', opts.title)
- ax.set_title(opts.title)
- if opts.xlabel:
- logging.debug('using x label: %r', opts.xlabel)
- ax.set_xlabel(opts.xlabel)
- if opts.xlim:
- logging.debug('using x limit: %r', opts.xlim)
- ax.set_xlim(eval(opts.xlim))
- if opts.ylabel:
- logging.debug('using y label: %r', opts.ylabel)
- ax.set_ylabel(opts.ylabel)
- if opts.ylim:
- logging.debug('using y limit: %r', opts.ylim)
- ax.set_ylim(eval(opts.ylim))
+ logging.debug('using grid: %s' % args.grid)
+ ax.grid(args.grid)
+ if args.title:
+ logging.debug('using title: %r', args.title)
+ ax.set_title(args.title)
+ if args.xlabel:
+ logging.debug('using x label: %r', args.xlabel)
+ ax.set_xlabel(args.xlabel)
+ if args.xlim:
+ logging.debug('using x limit: %r', args.xlim)
+ ax.set_xlim(eval(args.xlim))
+ if args.ylabel:
+ logging.debug('using y label: %r', args.ylabel)
+ ax.set_ylabel(args.ylabel)
+ if args.ylim:
+ logging.debug('using y limit: %r', args.ylim)
+ ax.set_ylim(eval(args.ylim))
-def main(opts, args):
- colors = [c.strip() for c in opts.colors.split(',')]
- points = [s.strip() for s in opts.points.split(',')]
- regex = compile_regex(opts)
- ax = make_axes(opts)
+def main(args):
+ regex = compile_regex(args)
+ ax = make_axes(args)
- plot_kwargs = dict(alpha=opts.alpha, aa=True)
+ plot_kwargs = dict(alpha=args.alpha, aa=True)
- c = p = 0
+ colors = itertools.cycle(args.colors)
+ points = itertools.cycle(args.points)
def plot(label, x, y, ey):
- if opts.smooth:
- y = numpy.convolve(y, [1. / opts.smooth] * opts.smooth, 'same')
+ if args.smooth:
+ y = numpy.convolve(y, [1. / args.smooth] * args.smooth, 'same')
- if opts.batch:
- n = opts.batch
+ if args.batch:
+ n = args.batch
count = int(numpy.ceil(float(len(y)) / n))
batches = lambda: (y[i * n:(i + 1) * n] for i in range(count))
means = [numpy.array(b).mean() for b in batches()]
stds = [numpy.array(b).std() for b in batches()]
y, ey = means, stds
- x = x or range(len(y))
+ x = x or list(range(len(y)))
- if opts.every:
- x = x[::opts.every]
- y = y[::opts.every]
+ if args.every:
+ x = x[::args.every]
+ y = y[::args.every]
- ax.plot(x, y, points[p], c=colors[c],
- mec=colors[c], mfc=(1, 1, 1, 1), mew=1.,
- label=label, **plot_kwargs)
+ ax.plot(x, y, next(points), c=next(colors), mew=0., label=label, **plot_kwargs)
if ey:
- if opts.every:
- ey = ey[::opts.every]
+ if args.every:
+ ey = ey[::args.every]
ax.errorbar(x, y, fmt=None, yerr=ey, ecolor=colors[c], **plot_kwargs)
- for label, lines in read_input(args):
+ for label, lines in read_input(args.input):
x, y, ey = [], [], []
for line in lines:
- search_line(line, regex, opts.column, x, y, ey)
+ search_line(line, regex, args.column, x, y, ey)
plot(label, x, y, ey)
- c = (c + 1) % len(colors)
- p = (p + 1) % len(points)
- format_axes(ax, opts)
+ format_axes(ax, args)
- if opts.output:
- logging.info('%s: saving plot', opts.output)
- return pyplot.savefig(opts.output)
+ if args.output:
+ logging.info('%s: saving plot', args.output)
+ return pyplot.savefig(args.output, dpi=args.dpi)
def on_close(event=None):
sys.exit()
@@ -308,4 +305,4 @@ if __name__ == '__main__':
stream=sys.stdout,
level=logging.INFO,
format='%(levelname).1s %(asctime)s %(message)s')
- main(*FLAGS.parse_args())
+ main(FLAGS.parse_args())
Please sign in to comment.
Something went wrong with that request. Please try again.