Use numba to speed up the loop

Question

From what I read, numba can significantly speed up the python program. Could my program time efficiency increase with numba?

import numpy as np def f_big(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100): return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2)) outer_sum = 0 dk = 0.000001 for k in np.arange(dk,0.4, dk): inner_sum = 0 for A in np.arange(dk, 20, dk): inner_sum += dk * f_big(A, k, 1e-5, 1e-5) outer_sum += inner_sum * dk print outer_sum 

Answer 1

Yes, this is the problem that Numba really works with. I changed the dk value because it was unreasonable for a simple demonstration. Here is the code:

 import numpy as np import numba as nb def f_big(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100): return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2)) def func(): outer_sum = 0 dk = 0.01 #0.000001 for k in np.arange(dk, 0.4, dk): inner_sum = 0 for A in np.arange(dk, 20, dk): inner_sum += dk * f_big(A, k, 1e-5, 1e-5) outer_sum += inner_sum * dk return outer_sum @nb.jit(nopython=True) def f_big_nb(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100): return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2)) @nb.jit(nopython=True) def func_nb(): outer_sum = 0 dk = 0.01 #0.000001 X = np.arange(dk, 0.4, dk) Y = np.arange(dk, 20, dk) for i in xrange(X.shape[0]): k = X[i] # faster to do lookup than iterate over an array directly inner_sum = 0 for j in xrange(Y.shape[0]): A = Y[j] inner_sum += dk * f_big_nb(A, k, 1e-5, 1e-5) outer_sum += inner_sum * dk return outer_sum 

And then the timings:

 In [7]: np.allclose(func(), func_nb()) Out[7]: True In [8]: %timeit func() 1 loops, best of 3: 222 ms per loop In [9]: %timeit func_nb() The slowest run took 419.10 times longer than the fastest. This could mean that an intermediate result is being cached 1000 loops, best of 3: 362 µs per loop 

So the numba version on my laptop is about 600 times faster.