We, data professionals, love to use NumPy for data calculation and data operations for arrays or matrices, however when it comes to 'for' loops, we might forget that NumPy can help us optimize away 'for' loops under certain circumstances with vectorization and broadcasting. In the process of doing that, we gain speed and performance. This article is to introduce three code patterns involving 'for' loops that could be replaced and optimized with NumPy operations.
Testing Data
Let's create a large enough test dataset with 30,000 elements for each variable.
for Loop with Item Check:
Below is a quick implementation of a single for loop with item check. WIth pure python, looping through a and checking if the item of a exists in b took about 376 ms.
Pure Python:
However, we could use np.isin(a, b) to get a boolean array and use the indexing functionality of np.array() to only return the elements with values of True. With this implementation, it only took 3.95ms to finish the operation, about 100x faster.
Numpy Optimization:
for Loop with if-else Condition:
Writing if-else conditions within a for loop is widely used. And below is a quick implementation of returning a string of “larger_than_10” if the item in a is larger than 10, or returning a string of “smaller_than_equal_10” if the item in a is smaller than or equal to 10.
Pure Python:
Numpy offers a convenient function, np.where() to handle if-else conditions and it took 525 microseconds to finish the operation, approximately 10x faster.
Numpy Optimization:
Double for Loops with Filter and Multiply Operations:
The pure python implementation below is to loop through all the items in a and loop through all the items in b and only multiple the item in a with the item in b when the item in a is larger than 10. Since this involves a double for loop, the operation was slow and it took about 2min 32s to finish.
Pure Python:
On the other hand, NumPy offers vectorized filter operation. With the help of reshaping the filtered array and broadcasting the multiply operation over the two arrays, we can replace the double for loop entirely with NumPy operations. And overall, it took only 539 ms to finish the operation and it's approximately 300x faster than the pure Python operation with double for loops.
NumPy Optimization:
Takeaway:
- NumPy operations are much faster than pure Python operations when you can find corresponding functions in NumPy to replace single for loops.
- Double for loops can sometimes be replaced by the NumPy broadcasting operation and it can save even more computational time.
Hope you find this helpful! Thanks for reading this week's tip! Check out other Data Hacks on Dataproducts.io.