bench.py

# -*- coding: utf-8 -*-
"""This module provides a simple class for benchmarking selected parts of the
code during regular execution of a program. Usage:
	with mark("category"):
		statements
will record the time and memory usage of statement each time it is
hit in the normal course of execution of the program. For each
category, the mean and standard deviation of time and memory
are computed. The statistics are available as the module-global
stats object. str(stats) will output something like this:

      n        time          cpu           mem           leak
            mean    std   mean    std   mean    std   mean    std
 foo  200   0.03  0.020   0.03  0.021   0.76   0.00   0.00  0.000
 bar 1010   0.00  0.015   0.00  0.016   0.76   0.00   0.00  0.000

The individual entries are also avilable by indexing stats like
a dictionary: stats["foo"]["time"].mean, .std, .n, .last
return the mean, standard deviation, number of hits and
last value for the execution time for category "foo", for example.
"""
import time, numpy as np
from collections import defaultdict
from enlib import memory

class Value:
	def __init__(self, n=0, v=0, vv=0):
		self.n  = n
		self.v  = v
		self.vv = vv
		self.last = v
	def add(self,v):
		self.n  += 1
		self.v  += v
		self.vv += v**2
		self.last = v
	@property
	def mean(self): return self.v/self.n if self.n > 0 else np.nan
	@property
	def std(self):
		try:
			return (self.vv/self.n - (self.v/self.n)**2)**0.5 if self.n > 0 else np.nan
		except ValueError:
			return np.nan
	def __repr__(self): return "Value(mean=%f, std=%f, n=%d)" % (self.mean, self.std, self.n)

class Entry(defaultdict):
	def __init__(self):
		defaultdict.__init__(self, Value)
	def __repr__(self):
		names = sorted(self)
		return "Entry("+", ".join(["%s=%s" % (name,str(self[name])) for name in names]) + ")"

class Register(defaultdict):
	def __init__(self, fmt=[("time","%6.2f","%6.3f",1),("cpu","%6.2f","%6.3f",1),("mem","%6.2f","%6.2f",2.0**30),("leak","%6.2f","%6.3f",2.0**30)]):
		defaultdict.__init__(self, Entry)
		self.info = fmt
	def add(self, name, *args):
		entry = self[name]
		for info,v in zip(self.info,args):
			entry[info[0]].add(v)
	def __repr__(self):
		# Sort by first ordering criterion
		if len(self) == 0: return ""
		nhit = np.array([self[name][self.info[0][0]].n    for name in self])
		vals = np.array([self[name][self.info[0][0]].mean for name in self])
		inds = np.argsort(nhit*vals)[::-1]
		names = list(self)
		names = [names[i] for i in inds]
		nhit  = nhit[inds]
		# For pretty output, determine ranges of each column
		name_dig = max([len(name) for name in names])
		nhit_dig = max([len("%d"%i) for i in nhit])
		lines = []
		# Header
		pre = " "*name_dig
		line1,line2 = pre + " %s" % "n".center(nhit_dig), pre + " "*(1+nhit_dig)
		for info in self.info:
			lens = [len(i%0) for i in info[1:3]]
			line1 += " " + info[0].center(sum(lens)+1)
			line2 += " " + " ".join("%*s" % (l,v) for l,v in zip(lens,["mean","std"]))
		lines.append(line1)
		lines.append(line2)
		for name,nh in zip(names,nhit):
			entry = self[name]
			line = "%-*s %*d" % (name_dig,name,nhit_dig,nh)
			for info in self.info:
				val = entry[info[0]]
				unit= float(info[3] if len(info) > 3 else 1)
				line += (" " + info[1] + " " + info[2]) % (val.mean/unit, val.std/unit)
			lines.append(line)
		return "\n".join(lines)
	def write(self, fname):
		with open(fname,"w") as f:
			f.write(str(self)+"\n")

stats = Register()

class mark:
	def __init__(self, name):
		self.name = name
	def __enter__(self):
		self.time1  = time.time()
		self.clock1 = time.clock()
		self.mem1   = memory.current()
	def __exit__(self, type, value, traceback):
		self.time2  = time.time()
		self.clock2 = time.clock()
		self.mem2   = memory.current()
		stats.add(self.name, self.time2 -self.time1, self.clock2-self.clock1, self.mem1, self.mem2-self.mem1)

class show:
	def __init__(self, name):
		self.name = name
	def __enter__(self):
		self.time1  = time.time()
		self.clock1 = time.clock()
		self.mem1   = memory.current()
	def __exit__(self, type, value, traceback):
		self.time2  = time.time()
		self.clock2 = time.clock()
		self.mem2   = memory.current()
		self.memmax = memory.max()
		print "%5.2f %5.2f %5.2f %s" % (self.time2-self.time1, self.mem2*1e-9, self.memmax*1e-9, self.name)