Mercurial > hg > numerics
changeset 53:cd9ec2784077
update plot package
author | Jeff Hammel <k0scist@gmail.com> |
---|---|
date | Tue, 20 Jan 2015 15:49:20 -0800 |
parents | 37838ae694d2 |
children | bbfe25d23a9f |
files | numerics/plot.py |
diffstat | 1 files changed, 146 insertions(+), 22 deletions(-) [+] |
line wrap: on
line diff
--- a/numerics/plot.py Mon Jan 19 13:24:43 2015 -0800 +++ b/numerics/plot.py Tue Jan 20 15:49:20 2015 -0800 @@ -2,44 +2,145 @@ # -*- coding: utf-8 -*- """ -plot data - -Usage:: +plot data with `matplotlib` - plot foo.csv - -Generates plots of all columns of foo versus the first column. -The title is foo.csv unless overridden with ``--title``. +See also : +- http://stackoverflow.com/questions/7534453/matplotlib-does-not-show-my-drawings-although-i-call-pyplot-show ; +- http://bokeh.pydata.org/ ; +- http://mpld3.github.io/ """ # imports import argparse +import csv +import matplotlib.cm as cm +import matplotlib.pyplot as plt +import numpy as np import os +import subprocess import sys +import tempfile +import time +from .utils import choose_program +from which import which +from StringIO import StringIO # module globals -__all__ = ['main', 'PlotParser'] +__all__ = ['Displayer', 'Plot', 'PlotParser', 'read', 'main'] string = (str, unicode) -def ensure_dir(directory): - """ensure a directory exists""" - if os.path.exists(directory): - assert os.path.isdir(directory) - return directory - os.makedirs(directory) - return directory + +class Displayer(object): + """image displayer""" + viewers = ('feh', 'qiv', 'sxiv') + + def __init__(self, viewers=None): + if viewers is None: + viewers = self.viewers + self.viewers = viewers + + def viewer(self): + """returns path to primary viewer""" + for viewer in self.viewers: + path = which(viewer) + if path: + return path + + def __call__(self, image): + """display an image and exit""" + viewer = self.viewer() + if viewer: + args = [viewer, image] + subprocess.check_output(args) + + +class Plot(object): + """plotting class""" + def __init__(self, title=None, xlabel=None, ylabel=None): + self.title = title + self.xlabel = xlabel + self.ylabel = ylabel + self._x = None + self._y = [] + self.marker = [] + + def x(self, data, label=None): + self._x = data + if label is not None: + self.xlabel = xlabel + + def y(self, data, label=None, marker='.'): + self._y.append(data) + self.marker.append(marker) + if label is not None: + self.ylabel = label + + def __call__(self, output): + assert self._y + + if self.title: + plt.title(self.title) + if self.xlabel: + plt.xlabel(self.xlabel) + if self.ylabel: + plt.ylabel(self.ylabel) + if self._x: + args = sum([[self._x, self._y[i], self.marker[i]] for i in range(len(self._y))], []) + plt.plot(*args) + else: + plt.plot(*self._y) + plt.show() + plt.savefig(output) + print ("{}->saved to '{}'".format(self.title or '', output)) + + +def read(f): + """ + Read from file ``f`` + Accepts CSV and space-delimited files + """ + + retval = None + for line in f: + line = line.strip() + if ',' in line: + buffer = StringIO() + buffer.write(line) + buffer.seek(0) + row = list(csv.reader(buffer))[0] + else: + row = line.split() + row = [float(i) for i in row] + if retval is None: + retval = [[i] for i in row] + else: + for index, value in enumerate(row): + retval[index].append(value) + + return retval class PlotParser(argparse.ArgumentParser): - """CLI option parser""" + """CLI option parser for the plotter""" - def __init__(self, **kwargs): + def __init__(self, *args, **kwargs): kwargs.setdefault('description', __doc__) - argparse.ArgumentParser.__init__(self, **kwargs) + argparse.ArgumentParser.__init__(self, *args, **kwargs) + self.add_argument('--info', dest='info', + help="display info and exit") + self.add_argument('input', + type=argparse.FileType('r'), nargs='*', + default=(sys.stdin,), + help='input file(s), or read from stdin if ommitted') + self.add_argument('-o', '--output', dest='output', + help="file name to output to") + self.add_argument('-s', '--scatter', dest='scatter', + action='store_true', default=False, + help="scatter plot") self.options = None - def parse_args(self, *args, **kw): - options = argparse.ArgumentParser.parse_args(self, *args, **kw) + def parse_args(self, *args, **kwargs): + options = argparse.ArgumentParser.parse_args(self, *args, **kwargs) self.validate(options) self.options = options return options @@ -51,11 +152,34 @@ def main(args=sys.argv[1:]): """CLI""" - # parse command line options + # parse command line parser = PlotParser() options = parser.parse_args(args) + plot_fcn = plt.scatter if options.scatter else plt.plot + + # read data + all_data = [read(f) for f in options.input] + + # color map + # http://stackoverflow.com/questions/12236566/setting-different-color-for-each-series-in-scatter-plot-on-matplotlib + n_col = sum([(len(data)-1) or 1 for data in all_data]) + colors = iter(cm.rainbow(np.linspace(0, 1, n_col))) + + for data in all_data: + # plot it + if len(data) == 1: + plot_fcn(*data, marker='.', color=next(colors)) + else: + for i in range(1, len(data)): + plot_fcn(data[0], data[i], label=str(i), marker='.', color=next(colors)) + + plt.show() + + # display + output = options.output or tempfile.mktemp(suffix='.png') + plt.savefig(output) + Displayer()(output) if __name__ == '__main__': main() -