# HG changeset patch
# User Jeff Hammel <k0scist@gmail.com>
# Date 1421797760 28800
# Node ID cd9ec2784077d61500d17d8930b63e0a4a661fb2
# Parent  37838ae694d28ba68b4139bf33733cb71106cf20
update plot package

diff -r 37838ae694d2 -r cd9ec2784077 numerics/plot.py
--- 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()
-