Optimized plot script to handle different directory structures.

1 files changed, 58 insertions, 96 deletions

diff --git a/plot_exps.py b/plot_exps.py
index 39529bd..bb6a707 100755
--- a/plot_exps.py
+++ b/plot_exps.py
@@ -1,15 +1,16 @@
 #!/usr/bin/env python
 from __future__ import print_function
 
+import matplotlib.pyplot as plot
 import os
 import shutil as sh
 import sys
+from collections import namedtuple
 from optparse import OptionParser
+from parse.col_map import ColMap
 from parse.dir_map import DirMap
 from parse.tuple_table import ReducedTupleTable
-from parse.col_map import ColMap
-from collections import namedtuple,defaultdict
-import matplotlib.pyplot as plot
+from plot.style import StyleMap
 
 def parse_args():
     parser = OptionParser("usage: %prog [options] [csv_dir]...")
@@ -21,15 +22,16 @@ def parse_args():
 
     return parser.parse_args()
 
-
 ExpDetails = namedtuple('ExpDetails', ['variable', 'value', 'title', 'out'])
 OUT_FORMAT = 'pdf'
 
-def get_details(path):
+def get_details(path, out_dir):
+    '''Decode a @path into details about a single experiment.'''
     out = "_".join(path) if path else "plot"
+    out = "%s/%s.%s" % (out_dir, out, OUT_FORMAT)
 
-    value = path.pop() if path else None
-    variable = path.pop() if path else None
+    value = path.pop(0) if path else None
+    variable = path.pop(0) if path else None
 
     title  = value.capitalize() if value else ""
     title += " by %s" % variable if variable else ""
@@ -37,105 +39,70 @@ def get_details(path):
 
     return ExpDetails(variable, value, title, out)
 
+def plot_by_variable(plot_node, col_map, details):
+    '''Plot each .csv files under @plot_node as a line on a shared plot.'''
 
+    # Generate mapping of (column)=>(line property to vary) for consistently
+    # formatted plots
+    columns = list(col_map.columns())
+    if details.variable and details.variable in columns:
+        columns.remove(details.variable)
+    style_map = StyleMap(columns, col_map.get_values())
 
-class StyleMap(object):
-    COLORS  = list('bgrcmyk')
-    LINES   = ['-', ':', '--']
-    MARKERS = list('.,ov^<>1234sp*hH+xDd|_')
-    ORDER   = [MARKERS, COLORS, LINES]
-    DEFAULT = ["k", "-", "k"]
-
-    def __init__(self, col_list, col_values):
-        self.prop_map = dict(zip(col_list, StyleMap.ORDER))
-
-        # Store 1 style per value
-        self.value_map = defaultdict(dict)
-        for column, styles in self.prop_map.iteritems():
-            value_styles = self.value_map[column]
-            for value in sorted(col_values[column]):
-                value_styles[value] = styles.pop(0)
-                styles += [value_styles[value]]
-
-    def get_style(self, kv):
-        style = ''
-        for k,v in kv.iteritems():
-            if k in self.value_map:
-                style += self.value_map[k][v]
-        return style
-
-    def get_key(self):
-        key = []
-        for column, properties in self.prop_map.iteritems():
-            idx = StyleMap.ORDER.index(properties)
-            prop_string = StyleMap.DEFAULT[idx] + "%s"
-            for value, prop in self.value_map[column].iteritems():
-                style = plot.plot([],[], prop_string%prop)[0]
-                key += [(style, "%s:%s" % (column, value))]
-        return sorted(key, key=lambda x:x[1])
-
-def plot_by_variable(dir_map, col_map, out_dir, force):
-    num_plots = 0
-    id = 0
-    for _,_ in dir_map.leafs(1):
-        num_plots += 1
-    sys.stderr.write("Plotting by variable...")
-
-    for plot_path, plot_node in dir_map.leafs(1):
-        id += 1
-        details = get_details(plot_path)
-        out_fname = "%s/%s.%s" % (out_dir, details.out, OUT_FORMAT)
-        if os.path.exists(out_fname) and not force:
-            continue
-
-        # Kinda bad...
-        first_csv = plot_node.children.keys()[0]
-        first_config = ColMap.decode(first_csv[:first_csv.index('.csv')])
-        columns = filter(lambda c: c in first_config, col_map.columns())
-
-        style_map = StyleMap(columns, col_map.get_values())
-
-        figure = plot.figure()
-        axes = figure.add_subplot(111)
+    figure = plot.figure()
+    axes = figure.add_subplot(111)
 
-        for line_path, line_node in plot_node.children.iteritems():
-            encoded = line_path[:line_path.index(".csv")]
-            config  = ColMap.decode(encoded)
-            style = style_map.get_style(config)
+    # Create a line for each file node
+    for line_path, line_node in plot_node.children.iteritems():
+        # Create line style to match this configuration
+        encoded = line_path[:line_path.index(".csv")]
+        config  = ColMap.decode(encoded)
+        style   = style_map.get_style(config)
 
-            values = sorted(line_node.values, key=lambda tup: tup[0])
-            xvalues, yvalues = zip(*values)
+        values = sorted(line_node.values, key=lambda tup: tup[0])
+        xvalues, yvalues = zip(*values)
 
-            plot.plot(xvalues, yvalues, style)
+        plot.plot(xvalues, yvalues, style.fmt())
 
-        lines, labels = zip(*style_map.get_key())
+    axes.set_title(details.title)
 
-        axes.legend(tuple(lines), tuple(labels), prop={'size':10})
-        axes.set_ylabel(details.value)
-        axes.set_xlabel(details.variable)
-        axes.set_xlim(0, axes.get_xlim()[1] + 1)
-        axes.set_ylim(0, axes.get_ylim()[1] + 1)
+    lines, labels = zip(*style_map.get_key())
+    axes.legend(tuple(lines), tuple(labels), prop={'size':10})
 
-        axes.set_title(details.title)
+    axes.set_ylabel(details.value)
+    axes.set_xlabel(details.variable)
+    axes.set_xlim(0, axes.get_xlim()[1] + 1)
+    axes.set_ylim(0, axes.get_ylim()[1] + 1)
 
-        plot.savefig(out_fname, format=OUT_FORMAT)
+    plot.savefig(details.out, format=OUT_FORMAT)
 
-        sys.stderr.write('\r {0:.2%}'.format(float(id)/num_plots))
-        sys.stderr.write('\n')
-
-def plot_exp(data_dir, out_dir, force):
-    print("Reading data...")
+def plot_dir(data_dir, out_dir, force):
+    sys.stderr.write("Reading data...\n")
     dir_map = DirMap.read(data_dir)
-    print("Sorting configs...")
+
+    sys.stderr.write("Creating column map...\n")
     tuple_table = ReducedTupleTable.from_dir_map(dir_map)
     col_map = tuple_table.get_col_map()
 
     if not os.path.exists(out_dir):
         os.mkdir(out_dir)
 
-    print("Plotting data...")
-    plot_by_variable(dir_map, col_map, out_dir, force)
-    # plot_by_config(tuple_table, out_dir)
+    sys.stderr.write("Plotting...\n")
+
+    # Count total plots for % counter
+    num_plots = len([x for x in dir_map.leafs(1)])
+    plot_num  = 0
+
+    for plot_path, plot_node in dir_map.leafs(1):
+        details = get_details(plot_path, out_dir)
+
+        if force or not os.path.exists(details.out):
+            plot_by_variable(plot_node, col_map, details)
+
+        plot_num += 1
+
+        sys.stderr.write('\r {0:.2%}'.format(float(plot_num)/num_plots))
+        sys.stderr.write('\n')
 
 def main():
     opts, args = parse_args()
@@ -146,13 +113,8 @@ def main():
     if not os.path.exists(opts.out_dir):
         os.mkdir(opts.out_dir)
 
-    for exp in args:
-        name = os.path.split(exp)[1]
-        if exp != os.getcwd():
-            out_dir = "%s/%s" % (opts.out_dir, name)
-        else:
-            out_dir = os.getcwd()
-        plot_exp(exp, out_dir, opts.force)
+    for dir in args:
+        plot_dir(dir, opts.out_dir, opts.force)
 
 if __name__ == '__main__':
     main()