Numba provides a Just In Time (JIT) compiler to generate optimized machine code to speed up computations in Python. It uses LLVM compiler infrastructure. One advantage of Numba over Cython is that there is no need to modify the Python code. Execution speed up is achieved using decorators, and if decorators are removed, we get back the pure Python code.
In [1]:
import numpy as np
def f(x, y):
return x + y # A trivial example
print f(1, 2) # int
print f(1.0, 2.0) # float
print f(1j, 2) # cmplx
print f([1, 2, 3], [4, 5, 6]) # list
print f(np.arange(1.0, 4.0), np.arange(4.0, 7.0)) # numpy.ndarray
In [2]:
a = np.array(np.arange(0, 100.1, 0.1), dtype=float)
print a.shape, a.size
%timeit a + a
%timeit f(a, a)
In [3]:
from numba import jit
@jit
def fast_f(x, y):
return x + y # A trivial example
%timeit fast_f(a, a)
In [5]:
from numba import f8, jit
def sum1d(x):
s = 0.0
for i in range(x.shape[0]):
s += x[i]
return s
@jit(f8(f8[:]))
def fast_sum1d(x):
s = 0.0
for i in range(x.shape[0]):
s += x[i]
return s
x = np.linspace(0, 100, 1001)
%timeit sum1d(x)
%timeit fast_sum1d(x)
Let us compute the approximate value of $\pi$ using the following sum:
$$\pi \approx \sqrt{6 \sum_{k=1}^{n+1} \frac{1}{k^2}}$$
In [6]:
def approx_pi(n=10000):
s = 0
for k in range(1, n+1):
s += 1.0 / k**2
return (6 * s) ** 0.5
%timeit approx_pi(1000)
%timeit approx_pi(10000)
print(approx_pi(1000))
print(approx_pi(10000))
In [7]:
from numba.decorators import autojit
@autojit
def fast_approx_pi(n):
s = 0
for k in range(1, n+1):
s += 1.0 / k**2
return (6 * s) ** 0.5
n = 10000
%timeit approx_pi(n)
%timeit fast_approx_pi(n)
In [8]:
from numba import vectorize, float64
@jit(float64(float64, float64))
def f(x, y):
return x + y
%timeit f(10, 20)
%timeit f(10.0, 20.0)
# ERROR - Will not work for arrays
a = np.array(np.arange(0, 100.1, 0.1), dtype=float)
%timeit f(a, a)
In [9]:
@jit(float64[:](float64[:], float64[:]))
def vf(x, y):
return x + y
%timeit vf(a, a)
%timeit vf(10, 20)
In [10]:
@vectorize([float64(float64, float64)])
def vfast_f(x, y):
return x + y # A trivial example
%timeit vfast_f(a, a)
%timeit vfast_f(10, 20)