Mercurial > hg > numerics
annotate numerics/plot.py @ 177:c2c86c37eadd
this now works on Mac
| author | Jeff Hammel <k0scist@gmail.com> |
|---|---|
| date | Thu, 10 Mar 2016 13:05:24 -0800 |
| parents | f62ac9a27a1f |
| children | 30d820087ae8 |
| rev | line source |
|---|---|
| 16 | 1 #!/usr/bin/env python |
| 2 # -*- coding: utf-8 -*- | |
| 3 | |
| 4 """ | |
| 53 | 5 plot data with `matplotlib` |
| 17 | 6 |
| 53 | 7 See also : |
| 8 - http://stackoverflow.com/questions/7534453/matplotlib-does-not-show-my-drawings-although-i-call-pyplot-show ; | |
| 9 - http://bokeh.pydata.org/ ; | |
| 10 - http://mpld3.github.io/ | |
| 119 | 11 |
| 12 For bokeh tools... | |
| 13 https://github.com/bokeh/bokeh/blob/master/bokeh/plotting_helpers.py#L277 | |
| 16 | 14 """ |
| 15 | |
| 16 # imports | |
| 17 import argparse | |
| 53 | 18 import csv |
| 19 import matplotlib.cm as cm | |
| 20 import matplotlib.pyplot as plt | |
| 21 import numpy as np | |
| 16 | 22 import os |
| 23 import sys | |
| 133 | 24 from .manipulate import ManipulationParser |
| 165 | 25 from StringIO import StringIO |
| 16 | 26 |
| 27 # module globals | |
|
129
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
28 __all__ = ['Plot', 'PlotParser', 'read', 'main'] |
| 16 | 29 string = (str, unicode) |
| 30 | |
| 53 | 31 |
| 32 class Plot(object): | |
| 33 """plotting class""" | |
| 34 def __init__(self, title=None, xlabel=None, ylabel=None): | |
| 35 self.title = title | |
| 36 self.xlabel = xlabel | |
| 37 self.ylabel = ylabel | |
| 38 self._x = None | |
| 39 self._y = [] | |
| 40 self.marker = [] | |
| 41 | |
| 42 def x(self, data, label=None): | |
| 43 self._x = data | |
| 44 if label is not None: | |
| 45 self.xlabel = xlabel | |
| 46 | |
| 47 def y(self, data, label=None, marker='.'): | |
| 48 self._y.append(data) | |
| 49 self.marker.append(marker) | |
| 50 if label is not None: | |
| 51 self.ylabel = label | |
| 52 | |
| 53 def __call__(self, output): | |
| 54 assert self._y | |
| 55 | |
| 56 if self.title: | |
| 57 plt.title(self.title) | |
| 58 if self.xlabel: | |
| 59 plt.xlabel(self.xlabel) | |
| 60 if self.ylabel: | |
| 61 plt.ylabel(self.ylabel) | |
| 62 if self._x: | |
| 63 args = sum([[self._x, self._y[i], self.marker[i]] for i in range(len(self._y))], []) | |
| 64 plt.plot(*args) | |
| 65 else: | |
| 66 plt.plot(*self._y) | |
| 67 plt.show() | |
| 68 plt.savefig(output) | |
| 69 print ("{}->saved to '{}'".format(self.title or '', output)) | |
| 70 | |
| 71 | |
| 72 def read(f): | |
| 73 """ | |
| 74 Read from file ``f`` | |
| 75 Accepts CSV and space-delimited files | |
| 76 """ | |
| 77 | |
| 177 | 78 if isinstance(f, string): |
| 79 with open(f, 'r') as _f: | |
| 80 return read(_f) | |
| 81 | |
| 53 | 82 retval = None |
| 83 for line in f: | |
| 84 line = line.strip() | |
| 85 if ',' in line: | |
| 86 buffer = StringIO() | |
| 87 buffer.write(line) | |
| 88 buffer.seek(0) | |
| 89 row = list(csv.reader(buffer))[0] | |
| 90 else: | |
| 91 row = line.split() | |
| 92 row = [float(i) for i in row] | |
| 93 if retval is None: | |
| 94 retval = [[i] for i in row] | |
| 95 else: | |
| 96 for index, value in enumerate(row): | |
| 97 retval[index].append(value) | |
| 98 | |
| 99 return retval | |
| 16 | 100 |
| 101 | |
| 133 | 102 class PlotParser(ManipulationParser): |
| 53 | 103 """CLI option parser for the plotter""" |
| 18 | 104 |
|
129
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
105 types = (float,) |
|
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
106 |
| 53 | 107 def __init__(self, *args, **kwargs): |
| 16 | 108 kwargs.setdefault('description', __doc__) |
| 133 | 109 ManipulationParser.__init__(self, *args, **kwargs) |
| 53 | 110 self.add_argument('-s', '--scatter', dest='scatter', |
| 111 action='store_true', default=False, | |
| 112 help="scatter plot") | |
| 16 | 113 self.options = None |
| 133 | 114 self.set_defaults(output=None) |
| 16 | 115 |
| 18 | 116 |
| 16 | 117 def main(args=sys.argv[1:]): |
| 118 """CLI""" | |
| 119 | |
| 53 | 120 # parse command line |
| 18 | 121 parser = PlotParser() |
| 16 | 122 options = parser.parse_args(args) |
| 123 | |
|
129
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
124 # choose plotting function |
| 53 | 125 plot_fcn = plt.scatter if options.scatter else plt.plot |
| 126 | |
| 127 # read data | |
| 169 | 128 all_data = [read(f) for f in options.input] |
| 53 | 129 |
| 130 # color map | |
| 131 # http://stackoverflow.com/questions/12236566/setting-different-color-for-each-series-in-scatter-plot-on-matplotlib | |
| 132 n_col = sum([(len(data)-1) or 1 for data in all_data]) | |
| 133 colors = iter(cm.rainbow(np.linspace(0, 1, n_col))) | |
| 134 | |
| 135 for data in all_data: | |
|
129
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
136 # plot each data set |
| 53 | 137 if len(data) == 1: |
| 138 plot_fcn(*data, marker='.', color=next(colors)) | |
| 139 else: | |
| 140 for i in range(1, len(data)): | |
| 141 plot_fcn(data[0], data[i], label=str(i), marker='.', color=next(colors)) | |
| 142 | |
|
121
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
143 # save plot |
|
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
144 if options.output: |
|
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
145 plt.savefig(options.output) |
| 165 | 146 print ("Figure saved to file://{}".format(os.path.abspath(options.output.name))) |
| 53 | 147 |
|
121
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
148 # display plot, I guess |
|
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
149 plt.show() |
| 18 | 150 |
| 177 | 151 |
| 16 | 152 if __name__ == '__main__': |
| 153 main() |