helper classes to generate graphs

benchmarks/graph.py

+# vim: ts=4:sw=4:et:ai:sts=4
+
+import csv, StringIO, subprocess
+
+class Graph:
+    [LINE, DOT, POINT, LINEPOINT] = range(0, 4)
+    def __init__(self):
+        self._plots = []
+        self._title = None
+    def set_title(self, title):
+        self._title = title
+    def add(self, plot):
+        self._plots.append(plot)
+    def generate(self, output_file):
+        lines = self.gen_output()
+        lines.insert(0, "set terminal postscript")
+        lines.insert(0, "set output '%s'" % filename)
+        gnuplot = subprocess.Popen(['gnuplot', '-'], 
+                stdin = subprocess.PIPE,
+                stdout = subprocess.PIPE,
+                stderr = subprocess.STDOUT)
+        gnuplot.communicate(input = "\n".join(lines))
+    def Xplot(self, plotnr):
+        lines = self.gen_output(plotnr)
+        lines.insert(0, "set terminal wxt")
+        lines.append('pause mouse')
+        gnuplot = subprocess.Popen(['gnuplot', '-'], stdin = subprocess.PIPE)
+        gnuplot.communicate(input = "\n".join(lines))
+    def _style_to_str(self, style):
+        d = {Graph.LINE: 'lines', Graph.DOT: 'dots', Graph.POINT: 'points', 
+                Graph.LINEPOINT: 'linespoints'}
+        return d[style]
+    def gen_output(self, plots = None):
+        if plots:
+            plots = map(self._plots.__getitem__, plots)
+        else:
+            plots = self._plots
+        lines = []
+        if self._title:
+            lines.append("set title '%s'" % self._title)
+        line = []
+        for plot in plots:
+            line.append("'-' title '%s' with %s" % (plot.title(),
+                self._style_to_str(plot.style())))
+        lines.append('plot ' + ', '.join(line))
+        for plot in plots:
+            for r in plot.data():
+                r = [str(d) for d in r]
+                lines.append(' '.join(r))
+            lines.extend(['e'])
+        return lines
+
+class Plot:
+    def __init__(self, title, data, style = Graph.LINE):
+        self._title = title
+        self._data = data
+        self._style = style
+    def style(self):
+        return self._style
+    def title(self):
+        return self._title
+    def data(self):
+        return self._data
+
+class Row:
+    def __init__(self, data, names = None):
+        assert not names or len(names) == len(data)
+        assert not names or all(map(lambda x: isinstance(x, str), names))
+        self._data1 = list(data)
+        if names:
+            self._data2 = dict(zip(names, data))
+        else:
+            self._data2 = dict()
+    def append(self, value, name = None):
+        self._data1.append(value)
+        if self._data2:
+            assert name not in self._data2
+            self._data2[name] = value
+    def __getitem__(self, item):
+        if isinstance(item, int):
+            return self._data1[item]
+        else:
+            return self._data2[item]
+    def __len__(self):
+        return len(self._data1)
+#    def __repr__(self):
+#        return 
+
+class Data:
+    def __init__(self, rows = [], colnames = []):
+        assert not (colnames and rows) or len(colnames) == len(rows[0])
+        self._colnames = colnames
+        self._data = []
+        for r in rows:
+            self.add_row(r)
+    def add_row(self, row):
+        if isinstance(row, Row):
+            self._data.append(row)
+        else:
+            self._data.append(Row(row, self._colnames))
+    def nrows(self):
+        return len(self._data)
+    def ncols(self):
+        return len(self._data[0])
+    def read_csv(self, stream, has_header = False):
+        self._data = []
+        self._colnames = []
+        self._datadict = []
+        n = 0
+        reader = csv.reader(stream)
+        for line in reader:
+            if n and len(line) != n:
+                raise 'Not matching number of columns in different rows'
+            if not n:
+                n = len(line)
+                if has_header:
+                    self._colnames = line
+                    continue
+            row = []
+            for i in line:
+                try:
+                    row.append(float(i))
+                except:
+                    row.append(i)
+            self._data.append(row)
+        if has_header:
+            self._gen_data_dict()
+    def write_csv(self, stream):
+        writer = csv.writer(stream)
+        writer.writerows(self._data)
+    def column(self, col):
+        if isinstance(col, int):
+            return [row[col] for row in self._data]
+        else:
+            return [row[col] for row in self._datadict]
+    def row(self, row):
+        return self._data[row]
+    def cell(self, row, col):
+        if isinstance(col, int):
+            return self._data[row][col]
+        else:
+            return self._datadict[row][col]
+    def select(self, cols = None, selectfn = lambda x: True,
+            groupfn = lambda x: None):
+        if cols:
+            cols = list(cols)
+        else:
+            cols = range(self.ncols())
+        groups = {}
+        res = []
+        for row in self._data if isinstance(cols[0], int) else self._datadict:
+            if not selectfn(row):
+                continue
+
+    def add_column(self, fn, colname = None):
+        if self._colnames:
+            assert colname
+        for row in self._data:
+            row.append(fn(row))
+        if colname:
+            self._colname.append(colname)
+            for row in self._datadict:
+                row[colname] = fn(row)
+        return self.ncols() - 1
+
+def uniq(l):
+    data = []
+    for i in l:
+        if i not in data:
+            data.append(i)
+    return data