1 | #!/usr/bin/env python |
---|
2 | |
---|
3 | '''A module to draw a graph on the screen given a data file.''' |
---|
4 | |
---|
5 | import sys |
---|
6 | import os.path |
---|
7 | import types |
---|
8 | import csv |
---|
9 | import time |
---|
10 | import getopt |
---|
11 | try: |
---|
12 | import pylab |
---|
13 | except: |
---|
14 | print 'Sorry, you have to install python-matplotlib' |
---|
15 | sys.exit(10) |
---|
16 | |
---|
17 | |
---|
18 | # Flag strings - keys in the 'options' dictionary |
---|
19 | X_DATACOL = 'x_datacol' |
---|
20 | Y_DATACOL = 'y_datacol' |
---|
21 | X_RANGE = 'x_range' |
---|
22 | Y_RANGE = 'y_range' |
---|
23 | FILENAME = 'filename' |
---|
24 | TITLE = 'title' |
---|
25 | X_LABEL = 'x_label' |
---|
26 | Y_LABEL = 'y_label' |
---|
27 | |
---|
28 | |
---|
29 | ## |
---|
30 | # @brief Issue an error message. |
---|
31 | # @param msg The message. |
---|
32 | def error(msg): |
---|
33 | print >>sys.stderr, msg |
---|
34 | sys.exit(10) |
---|
35 | |
---|
36 | |
---|
37 | ## |
---|
38 | # @brief Plot a sequence of data. |
---|
39 | # @param x_data The X data sequence to plot. |
---|
40 | # @param y_data A list of one or more Y data sequences to plot. |
---|
41 | # @param options A dictionary of plot options. |
---|
42 | def plot_data(x_data, y_data, options): |
---|
43 | pylab.plot(x_data, y_data) # x, y |
---|
44 | pylab.title(options.get(TITLE, '')) |
---|
45 | pylab.grid(True) |
---|
46 | |
---|
47 | # if user request a particular Y range |
---|
48 | if not options[Y_RANGE] is None: |
---|
49 | try: |
---|
50 | (minimum, maximum) = options[Y_RANGE].split(',') |
---|
51 | minimum = float(minimum) |
---|
52 | maximum = float(maximum) |
---|
53 | except: |
---|
54 | error('Sorry, got a bad value for Y range: %s' % options[Y_RANGE]) |
---|
55 | pylab.ylim(ymin=minimum, ymax=maximum) |
---|
56 | |
---|
57 | pylab.xlabel(options.get(X_LABEL, '')) |
---|
58 | pylab.ylabel(options.get(Y_LABEL, '')) |
---|
59 | |
---|
60 | pylab.show() |
---|
61 | |
---|
62 | |
---|
63 | ## |
---|
64 | # @brief Plot a data file. |
---|
65 | # @param filename Path to the data file to plot. |
---|
66 | # @param options A dictionary of options. |
---|
67 | def plot_file(filename, options=None): |
---|
68 | # get contents of data file |
---|
69 | # after this, 'header' is list of column header strings |
---|
70 | # 'data' is a list of lists of data |
---|
71 | fd = open(filename) |
---|
72 | c = csv.reader(fd) |
---|
73 | data = [] |
---|
74 | for row in c: |
---|
75 | data.append(row) |
---|
76 | fd.close() |
---|
77 | header = data[0] |
---|
78 | del data[0] |
---|
79 | |
---|
80 | # convert column specifiers to 'int' if required |
---|
81 | try: |
---|
82 | index = int(options[X_DATACOL]) |
---|
83 | except: |
---|
84 | try: |
---|
85 | index = header.index(options[X_DATACOL]) |
---|
86 | except ValueError: |
---|
87 | error("Sorry, X column header '%s' isn't in the data file." % options[X_DATACOL]) |
---|
88 | options[X_DATACOL] = index |
---|
89 | |
---|
90 | try: |
---|
91 | index = int(options[Y_DATACOL]) |
---|
92 | except: |
---|
93 | try: |
---|
94 | index = header.index(options[Y_DATACOL]) |
---|
95 | except ValueError: |
---|
96 | error("Sorry, Y column header '%s' isn't in the data file." % options[Y_DATACOL]) |
---|
97 | options[Y_DATACOL] = index |
---|
98 | |
---|
99 | # extract required columns from the data |
---|
100 | x_col = options[X_DATACOL] |
---|
101 | y_col = options[Y_DATACOL] |
---|
102 | |
---|
103 | # get max column number, check requested columns |
---|
104 | max_col = len(header) |
---|
105 | if x_col >= max_col or y_col >= max_col: |
---|
106 | error('Sorry, maximum column number for that file is %d.' % (max_col-1)) |
---|
107 | |
---|
108 | x_label = header[x_col].title() |
---|
109 | x_data = map(lambda x: x[x_col], data) |
---|
110 | if x_label == 'Time': |
---|
111 | x_label = 'Time (hours)' |
---|
112 | x_data = map(lambda x: float(x)/3600., x_data) |
---|
113 | y_data = map(lambda x: x[y_col], data) |
---|
114 | y_label = header[y_col].title() |
---|
115 | |
---|
116 | options[TITLE] = 'File: %s' % filename |
---|
117 | options[X_LABEL] = x_label |
---|
118 | options[Y_LABEL] = y_label |
---|
119 | |
---|
120 | plot_data(x_data, y_data, options) |
---|
121 | |
---|
122 | |
---|
123 | ## |
---|
124 | # @brief Help for the befuddled user. |
---|
125 | def usage(): |
---|
126 | print 'usage: %s <options> <filename>' % (ProgName) |
---|
127 | print 'where <filename> is the path to the data file to plot' |
---|
128 | print ' and <options> is zero or more of:' |
---|
129 | print ' -x <datacol> where <datacol> is a column specifier,' |
---|
130 | print ' -y <datacol> either a number or a header string' |
---|
131 | print ' -v <range> force a range in Y axis values, of the form' |
---|
132 | print ' <min>,<max>, eq, "-5,10"' |
---|
133 | print ' -h prints this help' |
---|
134 | sys.exit(10) |
---|
135 | |
---|
136 | if __name__ == '__main__': |
---|
137 | # dictionary to convert option string to internal name |
---|
138 | of_dict = { '-x': X_DATACOL, |
---|
139 | '-y': Y_DATACOL, |
---|
140 | '-v': Y_RANGE } |
---|
141 | |
---|
142 | # set name of the program (for error reporting) |
---|
143 | global ProgName |
---|
144 | try: |
---|
145 | ProgName = os.path.basename(sys.argv[0]) |
---|
146 | except: |
---|
147 | pass |
---|
148 | |
---|
149 | # get options |
---|
150 | opt_dict = {} |
---|
151 | params = sys.argv[1:] |
---|
152 | (opts, args) = getopt.gnu_getopt(params, 'x:y:v:h') |
---|
153 | if len(args) != 1: |
---|
154 | usage() |
---|
155 | for (o, v) in opts: |
---|
156 | try: |
---|
157 | opt_dict[of_dict[o]] = v.lower() |
---|
158 | except: |
---|
159 | usage() |
---|
160 | |
---|
161 | # get name of data file |
---|
162 | filename = args[0] |
---|
163 | (filename_minus, _) = filename.split('.', 1) |
---|
164 | |
---|
165 | # set default option values |
---|
166 | opt_dict[X_DATACOL] = opt_dict.get(X_DATACOL, 0) |
---|
167 | opt_dict[Y_DATACOL] = opt_dict.get(Y_DATACOL, 1) |
---|
168 | opt_dict[Y_RANGE] = opt_dict.get(Y_RANGE, None) |
---|
169 | opt_dict[FILENAME] = filename_minus |
---|
170 | |
---|
171 | # plot the file |
---|
172 | plot_file(filename, opt_dict) |
---|