How to time functions with timeit and save results

Viewed 84

I'm trying to compare the speed of my functions using timeit but saving the results of these and I can't achieve that. Let me show you a toy example, trying to compare diferent matrix multiplication functions:

from timeit import timeit
import numpy as np

def prod_tensorial (A,B)  : return np.tensordot(A,B,(-1,0))
def prod_punto (A,B)      : return np.dot(A,B)
def prod_sumeinsten (A,B) : return np.einsum('ij,jk->ik',A,B)

A = np.random.rand(1000,2000)
B = np.random.rand(2000,3000)

fs = ( prod_tensorial, prod_punto, prod_sumeinsten )
ts = [ ]
ps = [ ]

for f in fs:
    ts += [ timeit('ps += [f(A,B)]',number=1) ]     # error: ps & f didn't recognized
    #ts += [ timeit(lambda: f(A,B),number=1) ]      # it works but lossing results
    print( ts[-1] )

print(np.allclose(*ps))

How you can see, I can't save the results. Do you know how could I do that?

1 Answers

Function timeit does not have direct access to the global variables. You must provide the dictionary of the global variabes through the globals keyword. Since the variable ps is first read in the timed statement, it is by default considered local. Mark it as global:

# This part is only for testing
ps = []
def f(x,y): return x+y
A, B = 10, 11

timeit('global ps; ps += [f(A,B)]', number=1, globals=globals())
Related