tko/plotgraph.py

# Copyright Martin J. Bligh (mbligh@google.com), 2007

"""
Class to draw gnuplot graphs for autotest performance analysis.
Not that generic - specifically designed to to draw graphs of one type,
but probably adaptable.
"""

import subprocess, sys, os
from math import sqrt
Popen = subprocess.Popen

def avg_dev(values):
	if len(values) == 0:
		return (0,0)
	average = sum(values) / len(values)
	sum_sq_dev = sum( [(x - average) ** 2 for x in values] )
        std_dev = sqrt(sum_sq_dev / float(len(values)));
        return (average, std_dev);


class gnuplot:
	def __init__(self, title, xlabel, ylabel, xsort = sorted, size = "1180,900"):
		self.title = title
		self.xlabel = xlabel
		self.ylabel = ylabel
		self.data_titles = []
		self.datasets = []
		self.xsort = xsort
		self.xvalues = set([])
		self.size = size

	def xtics(self):
		count = 1
		tics = []
		for label in self.xsort(self.xlabels):
			tics.append('"%s" %d' % (label, count))
			count += 1
		return tics


	def add_dataset(self, title, labeled_values):
		"""
		Add a data line

		title: title of the dataset
		labeled_values: dictionary of lists
				{ label : [value1, value2, ... ] , ... }
		"""
		if not labeled_values:
			raise "plotgraph:add_dataset - dataset was empty! %s" %\
									title
		self.data_titles.append(title)
		data_points = {}
		for label in labeled_values:
			point = "%s %s" % avg_dev(labeled_values[label])
			data_points[label] = point
			self.xvalues.add(label)
		self.datasets.append(data_points)


	def plot(self, cgi_header = False, output = None, test = None):
		if cgi_header:
			print "Content-type: image/png\n"
			sys.stdout.flush()
		if test:
			g = open(test, 'w')
		else:
			p = Popen("/usr/bin/gnuplot", stdin = subprocess.PIPE)
			g = p.stdin
		g.write('set terminal png size %s\n' % self.size)
		g.write('set key below\n')
		g.write('set title "%s"\n' % self.title)
		g.write('set xlabel "%s"\n' % self.xlabel)
		g.write('set ylabel "%s"\n' % self.ylabel)
		if output:
			g.write('set output "%s"\n' % output)
		g.write('set style data yerrorlines\n')
		g.write('set grid\n')

		self.xlabels = self.xsort(list(self.xvalues))
		
		g.write('set xrange [0.5:%f]\n' % (len(self.xvalues)+0.5))
		g.write('set xtics rotate (%s)\n' % ','.join(self.xtics()))

		plot_lines = ['"-" title "%s"' % t for t in self.data_titles]
		g.write('plot ' + ', '.join(plot_lines) + '\n')

		for dataset in self.datasets:
			count = 1
			for label in self.xlabels:
				if not dataset.has_key(label):
					continue
				data = dataset[label]
				g.write("%d %s\n" % (count, str(data)))
				count += 1
			sys.stdout.flush()
			g.write('e\n')

		g.close()
		if not test:
			sts = os.waitpid(p.pid, 0)