The world’s leading publication for data science, AI, and ML professionals.

Say Goodbye to Loops and Hello to Optimization

Employing NumPy to Optimize Python Code

Photo by 𝓴𝓘𝓡𝓚 𝕝𝔸𝕀 on Unsplash
Photo by 𝓴𝓘𝓡𝓚 𝕝𝔸𝕀 on Unsplash

While Python remains easy to use, quick to learn, and offers an overabundance of external libraries that can do almost anything, it has one critical weakness: it’s slow.

Of course, to the human eye, its sluggishness seems negligible. Usually Python only lags behind other programming languages by milliseconds; however, when iterating over millions or even billions of data points, it quickly adds up.

NumPy offers a unique solution. While allowing users to still write Python, it converts it into well written C code specifically optimized for numerical analysis.

Using NumPy arrays can increase a script’s performance by one or two magnitudes with very little learning curve.

Vector Operations

Instead of performing an function on each element in an array, the traditional approach of loops, vector operations apply a function on all elements in an array simultaneously.

import random
# create a list of random integers
data = [random.randint(0, 100) for i in range(1000000)]

To demonstrate the utility of vector operations, an example dataset is created. The above code generates a list of 1,000,000 random integers between 0 and 100.

For the sake of example, the performance of taking the cosine of every element in the list will be measured.

# import math for cosine function
import math
# create a new list for the data
new_data = []
# loop through the data, take cosine and add it to the new list
for i in data:
    new_data.append(math.cos(i))

The above code uses a typical loop to go through the list and take the cosine of each value. This is the traditional approach for iteration that usually comes to mind. Unfortunately, this took about 0.31 seconds, which is rather slow.

# import math for cosine function
import math
# create a new list for the data
new_data = [math.cos(x) for x in data]

Instead of the loop, a list comprehension can also be used, which cuts down the time to 0.24 seconds. While an improvement over traditional loops, there’s still some sluggishness.

# transform the data into a NumPy array
data = numpy.array(data)
# apply the vector operation
new_data = numpy.cos(data)

Finally, NumPy is used for the same operation. The first line converts the list into an array and the second line applies the vector operation.

Note that instead of using the cosine function from the math module in the standard library, the cosine function from NumPy must be used instead.

In addition to generating cleaner and easier to read code, the entire process only takes about 0.03 seconds. This approximately 10 times faster than using a traditional loop and almost 90% faster than the list comprehension.

Filtering through Data

Another common task in which NumPy excels is data filtering. Using the same data as before, filtering for all values over 50 will be evaluated.

# Using a traditional loop
new_data = []
for i in data:
    if i > 50:
        new_data.append(i)
# Using a list comprehension
new_data = [i for i in data if i > 50]

As before, plain Python can filter through a list using either a traditional loop or a list comprehension. The former case required 0.23 seconds and the latter 0.11 seconds.

# Convert data into a NumPy array
data = numpy.array(data)
# Create a filter
data_filter = data > 50
# Filter for values greater than 50
new_data = data[data_filter]

Instead of using loops, however, NumPy offers an intuitive alternative. After converting the list into an array, the next line creates a filter. If iterated, the filter will output True and False values that correspond with condition.

In other words, if the first value in the data is greater than 50, the first value in the filter will be True and vice versa.

Once the filter is created, it can be applied to the data using familiar splicing notation, creating a new array with values greater than 50. The entire process only requires 0.02 seconds, a huge improvement over either loops or list comprehensions.

An Explanation of Broadcasting

Broadcasting concerns itself with arithmetic operations on arrays of unequal sizes. To work, the trailing axes of the arrays must be equal even if the number of dimensions aren’t. The illustration below gives a visual example.

A visual Example of Valid Broadcasting
A visual Example of Valid Broadcasting

In this case, broadcasting is compatible because there’s an equal number on the horizontal axis. The above example takes the 1 X 2 array and adds it to each row in the 2 X 2 array.

Another Visual Example of Valid Broadcasting
Another Visual Example of Valid Broadcasting

Alternatively, the vertical axis may be compatible. The same principle applies so that the 2 X 1 array is applied to each column of the 2 X 2 array.

Invalid Broadcasting
Invalid Broadcasting

If neither axis is equal, as in the example above, broadcasting may not be applied. Instead, Python returns a Value Error.

Code Examples of Broadcasting

Before applying broadcasting to higher dimensional data, the below example uses the same list of 1,000,000 integers between 0 and 100 used in previous demonstrations.

# Very simple broadcasting
new_data = numpy.array(data) * 2

Like before, instead of using a loop or list comprehension, simply converting the list into a NumPy array and applying the operation (multiplying by 2, in this case) provides a succinct and efficient way to transform the data.

From the perspective of broadcasting, a 1 X 1,000,000 array is multiplied by a 1 X 1 array (otherwise known as just a number).

Broadcasting, however, is far more interesting when applied to something more complex. Suppose the data is a 2 X 1,000,000 array instead:

# Create a 2 X 1,000,000 dataset
data = [
    [random.randint(0, 100) for i in range(1000000)],
    [random.randint(0, 100) for i in range(1000000)]
]

For the sake of example, everything in the first set of 1,000,000 data points will be added by 1 and add everything in the second set by 2.

new_data_1 = []
new_data_2 = []
# Using loops to inefficiently broadcast
for i in range(len(data)):
    for j in data[i]:
        new_data_1.append(data[i][j] + 1)

    for k in data[i]:
        new_data_2.append(data[i][k] + 2)

new_data = [new_data_1, new_data_2]

In a rather unfair comparison, a set of nested loops can be used to accomplished the goal, but the result is sloppy and needs a lot of overhead. It would also require a lot of tweaking if another set of dimensions were used.

Besides that, however, it takes 0.83 seconds to run.

# Using list comprehensions to broadcast
new_data = [
    [x + 1 for x in data[0]],
    [x + 2 for x in data[1]]
]

List comprehensions offer an attractive alternative. The above code is shorter, easier to read, and only takes 0.17 seconds to run; however, trying to use this method for more complex data (such as a 100 X 100 array) would be difficult to replicate.

# Convert data into a numpy array
data = numpy.array(data)
# Create a small array to broadcast
operator = numpy.array([
    [1],
    [2]
])
# Use NumPy broadcasting
new_data = data + operator

Finally, the example of broadcasting using NumPy provides a more intuitive approach. After converting the data into an array, a 2 X 1 array is created to modify the dataset. To accomplish this, a simple adding operator is used in the last line.

While more lines than using list comprehensions, this approach is more flexible, readable, and most importantly, only takes 0.04 seconds to run.

Conclusions

Python offers a lot of utility, but speed traditionally isn’t one of its benefits. The NumPy library, however, offers a large performance boost with an intuitive syntax.

Simply abandoning loops and list comprehensions in favor of NumPy operations can speed up code by a huge factor without adding an unneeded complexity. Understanding its use is a necessity for any large data project and practically a requirement for Data Science.


Related Articles