Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
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...
commit 74b56295f2ddbe44d4987f060b40a9fc4f8774c7 1 parent 1bec6c0
Leif Johnson authored
Showing with 121 additions and 50 deletions.
  1. +121 −50 scripts/py-grep-plot
171 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,13 +188,13 @@ 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:
@@ -162,27 +202,30 @@ def search_line(line, regex, columns, x, y, ey):
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,16 +291,22 @@ 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)))
@@ -261,25 +314,43 @@ def main(args):
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()
Please sign in to comment.
Something went wrong with that request. Please try again.