Permalink
Browse files

Move pyplot to plt and numpy to np. Add argument parsing from files. …

…Add options for auto-layout and for parsing dates on the x axis.
  • Loading branch information...
1 parent 1bec6c0 commit 74b56295f2ddbe44d4987f060b40a9fc4f8774c7 Leif Johnson committed Feb 19, 2013
Showing with 121 additions and 50 deletions.
  1. +121 −50 scripts/py-grep-plot
View
@@ -24,15 +24,18 @@
import argparse
import bz2
+import datetime
import glob
import gzip
+import itertools
import logging
-import numpy
+import numpy as np
import os
import re
import sys
-from matplotlib import pyplot
+from matplotlib import pyplot as plt
+from matplotlib import dates
LEGEND = {
'ul': 2, 'tl': 2,
@@ -46,24 +49,44 @@ LEGEND = {
'lr': 4, 'br': 4,
}
-FLAGS = argparse.ArgumentParser(
- conflict_handler='resolve',
- formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+class ArgParser(argparse.ArgumentParser):
+ SANE_DEFAULTS = dict(
+ fromfile_prefix_chars='@',
+ #conflict_handler='resolve',
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+ def __init__(self, *args, **kwargs):
+ kwargs.update(ArgParser.SANE_DEFAULTS)
+ super(ArgParser, self).__init__(*args, **kwargs)
+
+ def convert_arg_line_to_args(self, line):
+ '''Remove # comments and blank lines from arg files.'''
+ line = line.split('#')[0].rstrip()
+ if line:
+ if line[0] == '-' and ' ' in line:
+ for p in line.split():
+ yield p
+ else:
+ yield line
-FLAGS.add_argument('input', metavar='FILE', nargs=argparse.REMAINDER)
+FLAGS = ArgParser()
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('-A', '--auto', action='store_true',
+ help='layout plot automatically')
g.add_argument('-D', '--dpi', type=int, metavar='N',
help='save figure with dpi N')
+g.add_argument('-S', '--figsize', metavar='W,H',
+ help='save figure of size W x H inches')
+g.add_argument('-o', '--output', metavar='FILE',
+ help='save to FILE instead of displaying on screen')
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',
+g.add_argument('-k', '--column', nargs='*', 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')
@@ -77,10 +100,14 @@ 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.add_argument('-n', '--names', nargs='+', metavar='L',
+ help='use these names in the legend')
g = FLAGS.add_argument_group('axes')
g.add_argument('-g', '--grid', action='store_true',
help='include a grid')
+g.add_argument('-d', '--dates', metavar='FMT',
+ help='parse dates from x data using FMT')
g.add_argument('-L', '--legend', choices=tuple(sorted(LEGEND.keys())),
help='include a legend (None)')
g.add_argument('-l', '--log', choices=['x', 'y', 'xy'],
@@ -96,8 +123,11 @@ g.add_argument('-X', '--xlim', metavar='A,B',
g.add_argument('-Y', '--ylim', metavar='A,B',
help='use (A,B) as the range for the y-axis')
+FLAGS.add_argument('input', metavar='FILE', nargs=argparse.REMAINDER,
+ help='extract data for plotting from FILE')
+
-def extract_columns(data, columns, x, y, ey):
+def extract_columns(data, columns, x, y, ex, ey):
'''Pull specific column values out to plot.'''
if len(columns) == 1:
y.append(float(data[columns[0]]))
@@ -108,9 +138,14 @@ def extract_columns(data, columns, x, y, ey):
x.append(float(data[columns[0]]))
y.append(float(data[columns[1]]))
ey.append(float(data[columns[2]]))
+ if len(columns) == 4:
+ x.append(float(data[columns[0]]))
+ y.append(float(data[columns[1]]))
+ ey.append(float(data[columns[2]]))
+ ex.append(float(data[columns[3]]))
-def extract_groupdict(g, x, y, ey):
+def extract_groupdict(g, x, y, ex, ey):
'''We've matched a line with named groups. Extract data from them.'''
logging.debug('group dict: %r', g)
@@ -122,15 +157,20 @@ def extract_groupdict(g, x, y, ey):
y.append(float(g['y']))
if 'ey' in g:
- while len(ey) < len(y):
+ while len(ey) < len(y) - 1:
ey.append(None)
ey.append(float(g['ey']))
+ if 'ex' in g:
+ while len(ex) < len(x) - 1:
+ ex.append(None)
+ ex.append(float(g['ex']))
-def extract_groups(g, x, y, ey):
+
+def extract_groups(g, x, y, ex, ey):
logging.debug('group matches: %r', g)
if len(g) > 3:
- FLAGS.error('REGEX cannot match more than 3 values')
+ FLAGS.error('unnamed --regex cannot match more than 3 values')
elif len(g) == 3:
while len(x) < len(y):
x.append(None)
@@ -148,41 +188,44 @@ def extract_groups(g, x, y, ey):
y.append(float(g[0]))
-def search_line(line, regex, columns, x, y, ey):
+def search_line(line, regex, columns, *series):
'''Search an input line for groups matching the given regex.
- Extracted data will be added to the mutable x, y, and/or ey sequences.
+ Extracted data will be added to the mutable series sequences.
'''
- if columns:
- return extract_columns(line.split(), columns, x, y, ey)
+ if not regex:
+ return extract_columns(line.split(), columns, *series)
m = regex.search(line)
if not m:
return
g = m.groupdict()
if g:
- return extract_groupdict(g, x, y, ey)
+ return extract_groupdict(g, *series)
- extract_groups(m.groups(), x, y, ey)
+ extract_groups(m.groups(), *series)
-def read_input(args):
+def read_input(args, names):
'''Given input pattern arguments, open up corresponding files.'''
+ logging.debug('reading input data from %r', args)
if not args:
args.append('-')
- for pattern in args:
+ names = names or ()
+ for i, pattern in enumerate(args):
+ name = len(names) > i and names[i] or ''
if pattern == '-':
- yield 'stdin', sys.stdin
+ yield name or 'stdin', sys.stdin
continue
for filename in glob.glob(pattern):
if filename.endswith('.gz'):
- handle = gzip.open(filename)
+ handle = gzip.open(os.path.expanduser(filename))
elif filename.endswith('.bz2'):
- handle = bz2.BZ2File(filename)
+ handle = bz2.BZ2File(os.path.expanduser(filename))
else:
- handle = open(filename)
- yield os.path.splitext(os.path.basename(filename))[0], handle
+ handle = open(os.path.expanduser(filename))
+ yield name or os.path.splitext(os.path.basename(filename))[0], handle
def compile_regex(args):
@@ -197,13 +240,12 @@ def compile_regex(args):
def make_axes(args):
'''Create an axes object to hold our plots.'''
- 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 = plt.subplot(111)
ax.xaxis.tick_bottom()
ax.yaxis.tick_left()
- if 'x' in args.log:
+ if args.log and 'x' in args.log:
ax.set_xscale('log')
- if 'y' in args.log:
+ if args.log and 'y' in args.log:
ax.set_yscale('log')
return ax
@@ -233,6 +275,11 @@ def format_axes(ax, args):
if args.ylim:
logging.debug('using y limit: %r', args.ylim)
ax.set_ylim(eval(args.ylim))
+ if args.dates:
+ logging.debug('using dates from %r on x axis', args.dates)
+ loc = dates.AutoDateLocator()
+ ax.xaxis.set_major_locator(loc)
+ ax.xaxis.set_major_formatter(dates.AutoDateFormatter(loc))
def main(args):
@@ -244,42 +291,66 @@ def main(args):
colors = itertools.cycle(args.colors)
points = itertools.cycle(args.points)
- def plot(label, x, y, ey):
+ def plot(label, x, y, ex, ey):
+ if args.dates:
+ conv = lambda z: datetime.datetime.strptime(str(z), args.dates)
+ if args.dates == '%s':
+ conv = datetime.datetime.fromtimestamp
+ x = [dates.date2num(conv(z)) for z in x]
+
if args.smooth:
- y = numpy.convolve(y, [1. / args.smooth] * args.smooth, 'same')
+ y = np.convolve(y, [1. / args.smooth] * args.smooth, 'same')
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()]
+ count = int(np.ceil(float(len(y)) / n))
+ batches = [y[i * n:(i + 1) * n] for i in range(count)]
+ means = [np.array(b).mean() for b in batches]
+ stds = [np.array(b).std() for b in batches]
y, ey = means, stds
x = x or list(range(len(y)))
if args.every:
x = x[::args.every]
y = y[::args.every]
+ ex = ex[::args.every]
+ ey = ey[::args.every]
+
+ color = next(colors)
- ax.plot(x, y, next(points), c=next(colors), mew=0., label=label, **plot_kwargs)
+ if len(ex) and len(ey):
+ ex = np.asarray(ex)
+ ey = np.asarray(ey)
+ ax.errorbar(x, y, xerr=ex, yerr=ey, color=color, alpha=0.1, aa=True)
+ return
- if ey:
- if args.every:
- ey = ey[::args.every]
- ax.errorbar(x, y, fmt=None, yerr=ey, ecolor=colors[c], **plot_kwargs)
+ if len(ex):
+ ex = np.asarray(ex)
+ ax.fill_between(x + ex, x - ex, color=color, alpha=0.1, aa=True)
+ if len(ey):
+ ey = np.asarray(ey)
+ ax.fill_between(y + ey, y - ey, color=color, alpha=0.1, aa=True)
- for label, lines in read_input(args.input):
- x, y, ey = [], [], []
+ ax.plot(x, y, next(points), c=color, markeredgecolor=color, label=label, **plot_kwargs)
+
+ for label, lines in read_input(args.input, args.names):
+ series = [], [], [], []
for line in lines:
- search_line(line, regex, args.column, x, y, ey)
- plot(label, x, y, ey)
+ search_line(line, regex, args.column, *series)
+ plot(label, *series)
format_axes(ax, args)
+ if args.figsize:
+ plt.gcf().set_size_inches(*eval(args.figsize))
+
+ if args.auto:
+ plt.tight_layout()
+
if args.output:
logging.info('%s: saving plot', args.output)
- return pyplot.savefig(args.output, dpi=args.dpi)
+ return plt.savefig(os.path.expanduser(args.output), dpi=args.dpi)
def on_close(event=None):
sys.exit()
@@ -289,13 +360,13 @@ def main(args):
sys.exit()
try:
- pyplot.connect('close_event', on_close)
- pyplot.gcf().canvas.mpl_connect('key_press_event', on_key)
+ plt.connect('close_event', on_close)
+ plt.gcf().canvas.mpl_connect('key_press_event', on_key)
except ValueError:
pass
try:
- pyplot.show()
+ plt.show()
except KeyboardInterrupt:
quit()

0 comments on commit 74b5629

Please sign in to comment.