Module for conveniently comparing the performance of Python statements.

timestat.py

""" Python module for conveniently comparing the performance of
    different Python statements. 
    
    Features:
      - wraps `timeit`
      - one function call to compare multiple statements
      - textual and graphical representation of results
"""

import numpy as np
import matplotlib.pyplot as plt
import math
import random
import timeit
    
    
def compare_calls(stmts, setup='pass', globals=None, **kwargs):
    results = TimingResults()
    for s in stmts:
        n, t = time_stats(s, setup, globals, **kwargs)
        results.add(s, n, t)
    return results


def time_stats(stmt, setup='pass', globals=None, n_samples=100, burn_in_time=0.1, measure_time=0.9):
    
    timer = timeit.Timer(stmt, setup, globals=globals)
    
    start = timeit.time.time()
    target = start + burn_in_time
    n_loops, times, n = [], [], 1
    while timeit.time.time() < target:
        t = timer.timeit(n)
        times.append(t)
        n_loops.append(n)
        n *= 2
        
    total_loops = sum(n_loops)
    total_time = sum(times)
    
    print(total_loops, total_time)
    
    x = 2 * measure_time * total_loops / (total_time * n_samples * (n_samples + 1))
    
    n_loops = [math.ceil(x * n) for n in range(1, n_samples+1)]
    random.shuffle(n_loops)

    start = timeit.time.time()
    results = []
    for n in n_loops:
        t = timer.timeit(n)
        results.append((n, t))
    actual = timeit.time.time() - start

    print(actual)

    n_loops = np.array([n for (n, t) in results])
    times = np.array([t for (n, t) in results])
    
    return n_loops, times


class TimingResults:
    def __init__(self):
        self.results = []
    
    def add(self, name, n, t):
        self.results.append((name, n, t))
        
    def hist(self, ax=None):
        if ax is None:
            fig = plt.figure()
            ax = plt.subplot(1, 1, 1)
        for s, n, t in self.results:
            ax.hist(t/n, int(np.sqrt(len(n))), label=s, alpha=0.8)
        plt.semilogx()
        plt.xlabel('seconds per call')
        plt.legend()
        
    def slope(self, ax=None):
        if ax is None:
            fig = plt.figure()
            ax = plt.subplot(1, 1, 1)
        for s, n, t in self.results:
            ax.scatter(n, t, alpha=1 / np.log10(len(n)), label=s)
        plt.xlabel('repetitions')
        plt.ylabel('time')
        plt.legend()
        
    def __str__(self):
        if len(self.results) == 0:
            return 'no results...'   

        ml = max(len(name) for name, _, _ in self.results)
        ml = min(ml, 30)
        
        res = []
        for name, n, t in self.results:        
            if len(name) > ml:
                res += [name, ' :\n', ' ' * ml, '   ']
            else:
                res += [' ' * (ml - len(name)), name, ' : ']
            res += ['{:2.2} s/call median, {:2.2} ... {:2.2} IQR\n'.format(np.median(t/n), np.percentile(t/n, 25), np.percentile(t/n, 75))]
        return ''.join(res)
    
    
if __name__ == '__main__':

    def fib_slow(n):
        if n < 2:
            return 1
        return fib_slow(n-1) + fib_slow(n-2)
        
    def fib_fast(n):
        a, b = 1, 1
        for i in range(n):
            a, b = b, a + b
        return a
        
    def fib_cheat(n):
        return [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233]
            
    result = compare_calls(['fib_slow(10)', 'fib_fast(10)', 'fib_cheat(10)'], globals=globals())
    
    print('\n==========\n')
    
    print(result)
    result.slope()
    result.hist()
    
    plt.show()