Credits : Unsplash

Being efficient with Numpy Array Indexing — Part 1

This article explains an example of how to use numpy indexing efficiently.

Let’s start with a simple example of transforming elements of an array by an arbitrary transformation rule.

Say you have a 1D array of integers named source which you want to transform into 1D array of integers named result using the lookup table named lookup_table.

source = [0,2,3,1,4,5]
lookup_table = [11,22,33,44,55,66]

How to perform lookup?

To find the mapping of an integer n read nth element in lookup_table.

If we apply above transformation rule to source array, we would get

result = [11,33,44,22,55,66]

A naive python code for doing this would look like this:

import time
source = [0,2,3,1,4,5]
lookup_table = [11,22,33,44,55,66]
result = []
start = time.time()
for element in source:
    transformed_value = lookup_table[element]
    result.append(transformed_value)
end = time.time()
print(result) # [11,33,44,22,55,66]
print(end - start) 

This approach works well with smaller list/array of data. Using same approach for a list / array of 10 million elements will take significant cpu time.

What must be slowing the above code?

As you guessed correctly, it is for loop. It is well known that looping in python is damn slow due to interpreted nature of Python. If you want to perform same operations with a 2D array, performance will further degrade due to nested for loops. That is where Numpy comes for the rescue.

NumPy is the fundamental package for scientific computing with Python.

A vectorized implementation of above code will be hundred times faster in Numpy. Let’s see how to do it.

import numpy as np
import time
source = np.array([0,2,3,1,4,5])
lookup_table = np.array([11,22,33,44,55,66])
result = np.zeros_like(source)
start = time.time()
result = lookup_table.ravel()[source.ravel()]
end = time.time()
print(end -start)

Let me explain.

Time being ignore the .ravel() part. Then what you see is

result = lookup_table[source]

We can use a numpy array to index another numpy array and that is what I’m doing here. We are using elements of numpy to index lookuptable. What you see above is sometimes referred to as vectorized implementation. And vectorized implementations are always faster than loops( at least in Python!).

Now the .ravel() part. Doing .ravel() on a numpy array is as good as reshaping the array to a column vector. We are doing it so that we can index any element of lookup_table with only one index.

Let’s see if the vectorised version is faster than loop version. Let’s use above code snippets with a million elements.

For the first snippet, it becomes,

import time
source = [0,2,3,1,4,5] * 1000000
lookup_table = [11,22,33,44,55,66]
result = []
start = time.time()
for element in source:
    transformed_value = lookup_table[element]
    result.append(transformed_value)
end = time.time()
print(result[:5]) # [11,33,44,22,55]
print(end - start) # ~2.5 seconds

And for the second snippet,

import numpy as np
import time
source = np.array([0,2,3,1,4,5] * 1000000)
lookup_table = np.array([11,22,33,44,55,66])
result = np.zeros_like(source)
start = time.time()
result = lookup_table.ravel()[source.ravel()]
end = time.time()
print(end -start) # ~0.026 seconds

As you can see, vectorised version is approximately 100 times faster than the loop version.

How to deal with mapping problems in 2D?

Simple, do .ravel() on both source and lookup_table and then resize the result back to it’s original dimensions.

rows, columns = source.shape[:2]
result = lookup_table.ravel()[source.ravel()]
result = result.reshape((rows, columns))

Thanks!