Add --help, --stack, and --output to sawsim_plot_histogram_file.py.
authorW. Trevor King <wking@tremily.us>
Fri, 27 Jul 2012 20:23:41 +0000 (16:23 -0400)
committerW. Trevor King <wking@tremily.us>
Fri, 27 Jul 2012 20:23:41 +0000 (16:23 -0400)
bin/sawsim_plot_histogram_file.py

index 25e2f96..9501d2d 100755 (executable)
@@ -1,15 +1,62 @@
 #!/usr/bin/env python
 
+"""Plot a sawsim histogram file.
+"""
+
+import argparse as _argparse
+
+from matplotlib import pyplot
+
 from pysawsim.parameter_scan import HistogramMatcher
+from pysawsim.histogram import FIGURE
 
 
-if __name__ == '__main__':
-    import sys
-    
-    for histfile in sys.argv[1:]:
+def main(argv=None):
+    parser = _argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('files', metavar='FILE', nargs='+')
+    parser.add_argument(
+        '-s', '--stack', action='store_const', const=True,
+        help='stack all histograms into a single plot')
+    parser.add_argument(
+        '-o', '--output',
+        help='filename for saving the stacked output plot')
+
+    args = parser.parse_args(argv)
+
+    if args.output and not args.stack:
+        raise ValueError('cannot set --output without setting --stack')
+
+    stack = []
+    for histfile in args.files:
         with open(histfile, 'r') as f:
             histograms = HistogramMatcher._read_force_histograms(f)
             for params,histogram in sorted(histograms.items()):
-                histogram.plot(
-                    title=params,
-                    filename='histogram{}.png'.format(params.replace(' ','_')))
+                if args.stack:
+                    stack.append((params,histogram))
+                else:
+                    histogram.plot(
+                        title=params,
+                        filename='histogram{}.png'.format(
+                            params.replace(' ','_')))
+    if stack:
+        # fairly similar to pysawsim.histogram.Histogram.plot
+        FIGURE.clear()
+        axes = FIGURE.add_subplot(1, 1, 1)
+        axes.hold(True)
+        bin_width = stack[0][1].bin_edges[1] - stack[0][1].bin_edges[0]
+        dx = bin_width / (2*len(stack))  # small offset for readability
+        for i,(params,histogram) in enumerate(stack):
+            n,bins,patches = axes.hist(
+                x=[x + i*dx for x in histogram.bin_edges[:-1]],
+                weights=histogram.probabilities,
+                bins=[x + i*dx for x in histogram.bin_edges],
+                align='mid', histtype='stepfilled',
+                alpha=0.15, label=params)
+            axes.legend(loc='best')
+        axes.set_title('histograms')
+        pyplot.show()
+        FIGURE.savefig(args.output)
+
+
+if __name__ == '__main__':
+    main()