Mercurial > hg > numerics
view numerics/plot.py @ 63:0df8bcb6d521
stubbing: unicode histograms
author | Jeff Hammel <k0scist@gmail.com> |
---|---|
date | Thu, 26 Feb 2015 15:12:48 -0800 |
parents | cd9ec2784077 |
children | 3602e357d5e7 |
line wrap: on
line source
#!/usr/bin/env python # -*- coding: utf-8 -*- """ plot data with `matplotlib` 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__ = ['Displayer', 'Plot', 'PlotParser', 'read', 'main'] string = (str, unicode) 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 for the plotter""" def __init__(self, *args, **kwargs): kwargs.setdefault('description', __doc__) 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, **kwargs): options = argparse.ArgumentParser.parse_args(self, *args, **kwargs) self.validate(options) self.options = options return options def validate(self, options): """validate options""" def main(args=sys.argv[1:]): """CLI""" # 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()