5.4.1 Vectorization

To improve performance, you have often to vectorize the code. Replacing for loops and other slower parts of the code with NumPy slicing, operations, and functions can give significant improvements.

For example, the simple addition of a scalar to a vector by iterating over the elements is very slow:

for i in range(len(v)):
    w[i] = v[i] + 5

But using NumPy's addition is much faster:

w = v + 5

Using NumPy slicing can also give significant speed improvements over iterating with for loops. To demonstrate this, let's consider forming the average of neighbors in a two-dimensional array:

def my_avg(A):
    m,n = A.shape
    B = A.copy()
    for i in range(1,m-1):
        for j in range(1,n-1):
            B[i,j] = (A[i-1,j] + A[i+1,j] + A[i,j-1] + A[i,j+1])/4
    return B

def slicing_avg(A):
    A[1:-1,1:-1] = (A[:-2,1:-1] + A[2:,1:-1] +
                    A[1:-1,:-2] + A[1:-1,2:])/4
    return A

These functions both assign each element...