Extend src/sorting/scaling.py to handle `cat data | sort` syntax.
[parallel_computing.git] / src / sorting / scaling.py
1 #!/usr/bin/env python
2
3 """Measure how a sorting executable scales with N.
4
5 The executable should support one of the following:
6   executable path/to/data/file
7   cat path/to/data/file | executable
8 Where the data file is of the format output by data.py.
9 """
10
11 import subprocess
12 import time
13 import sys
14
15 import numpy
16 import matplotlib
17 matplotlib.use('Agg')  # select backend that doesn't require X Windows
18 import pylab
19
20
21 def generate_data(stream, N, ordered=False):
22     print >> sys.stderr, 'generate %d data points (ordered? %s)' % (
23         N, ordered)
24     stream.seek(0)
25     stream.truncate()
26     args = ['./data.py', str(N)]
27     if ordered:
28         args.insert(1, '--ordered')
29     q = subprocess.Popen(args, stdout=stream)
30     status = q.wait()
31     assert status == 0, status
32     stream.flush()
33
34 def run_test(executable, stdin=True, data_filename=None):
35     print >> sys.stderr, 'run %s' % executable
36     if stdin:
37         with open(data_filename, 'r') as f:
38             contents = f.read()
39     start = time.time()
40     if stdin:
41         p = subprocess.Popen([executable],
42                              stdout=open('/dev/null', 'w'))
43         p.communicate(contents)
44     else:
45         p = subprocess.Popen([executable, data_filename],
46                              stdout=open('/dev/null', 'w'))
47     status = p.wait()
48     stop = time.time()
49     assert status == 0, status
50     return stop - start
51     
52 def run_tests(executable, stdin=True, data_file=None, ordered=False,
53               repeats=10, max_time=1e2):
54     times = {}
55     prev_time = 0
56     N = 2
57     while prev_time < max_time:
58         print 
59         ts = numpy.zeros((repeats,), dtype=numpy.double)
60         for i in range(repeats):
61             generate_data(data_file, N, ordered=ordered)
62             ts[i] = run_test(executable, stdin, data_file.name)
63         times[N] = ts
64         prev_time = ts.mean()
65         N *= 2
66     return times
67
68
69 if __name__ == '__main__':
70     import optparse
71     import tempfile
72
73     p = optparse.OptionParser(
74         usage='%prog [options] executable', epilog=__doc__)
75     p.add_option('-s', '--stdin', dest='stdin', default=False,
76                  action='store_true', help='Use the stdin executable syntax.')
77     p.add_option('-r', '--repeats', dest='repeats', default=10, type='int',
78                  help='Number of repeats to run at each N (%default).')
79     p.add_option('-m', '--max-time', dest='max_time', default=1e2,type='float',
80                  help='Number of repeats to run at each N (%default).')
81     p.add_option('-p', '--plot', dest='plot', default=None,
82                  help='Filename for a scaling plot (no plot is generated if this option is not set).')
83
84     options,args = p.parse_args()
85
86     executable = args[0]
87
88     data_file = tempfile.NamedTemporaryFile()
89     kwargs = {
90         'executable': executable,
91         'stdin': options.stdin,
92         'data_file': data_file,
93         'repeats': options.repeats,
94         'max_time': options.max_time,
95         }
96     try:
97         times = run_tests(ordered=False, **kwargs)
98         ordered_times = run_tests(ordered=True, **kwargs)
99     except:
100         data_file.close()
101         raise
102
103     columns = ['N',
104                'ordered mean (s)', 'ordered std. dev. (s)',
105                'random mean (s)', 'random std. dev. (s)']
106     plots = dict([(c, []) for c in columns])
107
108     print '# sort times for %s' % executable
109     print '# %d repeats' % options.repeats
110     print '#%s' % '\t'.join(columns)
111     invalid = numpy.array(numpy.inf)
112     for key in sorted(set(times.keys() + ordered_times.keys())):
113         om = ordered_times.get(key, invalid).mean()
114         os = ordered_times.get(key, invalid).std()
115         m = times.get(key, invalid).mean()
116         s = times.get(key, invalid).std()
117         print '\t'.join([str(x) for x in [key, om, os, m, s]])
118         for c,x in zip(columns, [key, om, os, m, s]):
119             plots[c].append(x)
120
121     if options.plot:
122         f = pylab.figure()
123         a = pylab.axes()
124         a.hold(True)
125         for c,color in zip(['ordered', 'random'], 'br'):
126             a.errorbar(
127                 x=plots['N'],
128                 y=plots['%s mean (s)' % c],
129                 yerr=plots['%s std. dev. (s)' % c],
130                 fmt='%so-' % color, label=c)
131         a.set_title('sort times for %s' % executable)
132         a.set_xscale('log')
133         a.set_yscale('log')
134         a.set_xlabel('N')
135         a.set_ylabel('t (s)')
136         a.legend(loc='best')
137         f.savefig(options.plot)