Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Use numba to speed up for loop

Tags:

python

numba

From what I've read, numba can significantly speed up a python program. Could my program's time efficiency be increased using 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
like image 596
kilojoules Avatar asked Mar 03 '16 21:03

kilojoules


People also ask

What can Numba speed up?

For certain types of computation, in particular array-focused code, the Numba library can significantly speed up your code. Sometimes you'll need to tweak it a bit, sometimes it'll just work with no changes. And when it works, it's a very transparent speed fix.

How do you speed up a loop in Python?

A faster way to loop using built-in functions A faster way to loop in Python is using built-in functions. In our example, we could replace the for loop with the sum function. This function will sum the values inside the range of numbers. The code above takes 0.84 seconds.

Does Numba make NumPy faster?

Introduction. For the uninitiated Numba is an open-source JIT compiler that translates a subset of Python/NumPy code into an optimized machine code using the LLVM compiler library. In short Numba makes Python/NumPy code runs faster. It achieves this by compiling your Python code into native machine code.


1 Answers

Yes, this is the sort of problem that Numba really works for. I changed your value of dk because it wasn't sensible 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 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 is approx 600 times faster on my laptop.

like image 51
JoshAdel Avatar answered Sep 21 '22 12:09

JoshAdel