Mercurial > hg > numerics
annotate numerics/plot.py @ 172:dc0a620a0368
add another example
author | Jeff Hammel <k0scist@gmail.com> |
---|---|
date | Fri, 03 Jul 2015 11:11:01 -0700 |
parents | f62ac9a27a1f |
children | c2c86c37eadd |
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 | |
78 retval = None | |
79 for line in f: | |
80 line = line.strip() | |
81 if ',' in line: | |
82 buffer = StringIO() | |
83 buffer.write(line) | |
84 buffer.seek(0) | |
85 row = list(csv.reader(buffer))[0] | |
86 else: | |
87 row = line.split() | |
88 row = [float(i) for i in row] | |
89 if retval is None: | |
90 retval = [[i] for i in row] | |
91 else: | |
92 for index, value in enumerate(row): | |
93 retval[index].append(value) | |
94 | |
95 return retval | |
16 | 96 |
97 | |
133 | 98 class PlotParser(ManipulationParser): |
53 | 99 """CLI option parser for the plotter""" |
18 | 100 |
129
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
101 types = (float,) |
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
102 |
53 | 103 def __init__(self, *args, **kwargs): |
16 | 104 kwargs.setdefault('description', __doc__) |
133 | 105 ManipulationParser.__init__(self, *args, **kwargs) |
53 | 106 self.add_argument('-s', '--scatter', dest='scatter', |
107 action='store_true', default=False, | |
108 help="scatter plot") | |
16 | 109 self.options = None |
133 | 110 self.set_defaults(output=None) |
16 | 111 |
18 | 112 |
16 | 113 def main(args=sys.argv[1:]): |
114 """CLI""" | |
115 | |
53 | 116 # parse command line |
18 | 117 parser = PlotParser() |
16 | 118 options = parser.parse_args(args) |
119 | |
129
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
120 # choose plotting function |
53 | 121 plot_fcn = plt.scatter if options.scatter else plt.plot |
122 | |
123 # read data | |
169 | 124 all_data = [read(f) for f in options.input] |
53 | 125 |
126 # color map | |
127 # http://stackoverflow.com/questions/12236566/setting-different-color-for-each-series-in-scatter-plot-on-matplotlib | |
128 n_col = sum([(len(data)-1) or 1 for data in all_data]) | |
129 colors = iter(cm.rainbow(np.linspace(0, 1, n_col))) | |
130 | |
131 for data in all_data: | |
129
79d99216ba17
cleanup and remove image display
Jeff Hammel <k0scist@gmail.com>
parents:
121
diff
changeset
|
132 # plot each data set |
53 | 133 if len(data) == 1: |
134 plot_fcn(*data, marker='.', color=next(colors)) | |
135 else: | |
136 for i in range(1, len(data)): | |
137 plot_fcn(data[0], data[i], label=str(i), marker='.', color=next(colors)) | |
138 | |
121
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
139 # save plot |
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
140 if options.output: |
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
141 plt.savefig(options.output) |
165 | 142 print ("Figure saved to file://{}".format(os.path.abspath(options.output.name))) |
53 | 143 |
121
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
144 # display plot, I guess |
e259aa9ee18f
fix up how this works slightly
Jeff Hammel <k0scist@gmail.com>
parents:
120
diff
changeset
|
145 plt.show() |
18 | 146 |
16 | 147 if __name__ == '__main__': |
148 main() |