In this article, we’ll learn Numpy Vectorization in Python. Numpy is a C implementation of arrays in Python that is comparatively faster while having the same Python interpreter. In this article, we explore how we can apply a function element-wise to a sequence using numpy.vectorize().
Why for loop isn’t best for Numpy Vectorization?
To apply a function to an element of a list or tuple or a NumPy array, we can easily use the for loop in Python. But Python is an interpreted language and most of the implementation is slow compared to that of C and C++. The main reason for this slow computation comes down to the dynamic nature of Python and the lack of compiler level optimizations which incur memory overheads.
This is not an ideal situation for people who use Python for huge computations. Though NumPy provides faster implementation, for loop takes away some of that speed NumPy offers. To tackle this bottleneck, NumPy provides vectorization functionality that maps a function over a sequence efficiently.
numpy.vectorize() vs Python for loop – Vectorization speed comparison
So let us the test the speed of the python for loop vs the vectorized version. We’ll use the timeit function to get an accurate speed test.
# We use a large array for benchmarking our method
a = np.random.rand(10000)
b = 5
print("Benchmark for the for loop implementation: ")
%timeit [foo(i, b) for i in a]
print()
print("Benchmark for the vecfoo implementation: ")
%timeit vecfoo(a, b)
We see that the vectorized version is more than 3 times faster than the for loop implementation.
Numpy Vectorization with the numpy.vectorize() function
Numpy vectorize function takes in a python function (pyfunc) and returns a vectorized version of the function.
The vectorized version of the function takes a sequence of objects or NumPy arrays as input and evaluates the Python function over each element of the input sequence. Numpy Vectorization essentially functions like the python map() but with additional functionality – the NumPy broadcasting mechanism.
So let us understand the the numpy.vectorize() function in more detail:
numpy.vectorize(pyfunc, otypes=None, doc=None, excluded=None, cache=False, signature=None)
Required Parameters:
pyfunc: The function that we want to apply over the sequence of objects
Optional Parameters:
- otypes: The output types of the function can be specified as a string or a list of datatypes. If the otypes are not mentioned and the cache is set to True, the output type is determined by calling the first element of the input.
- doc: To specify the docstring of the created. If not mentioned the original docstring of the function( if any) will be used.
- cache: If True, then cache the first function call that determines the number of outputs if otypes is not provided.
Vectorizing a function
def foo(a, b):
"""
If a > b return a + b,
else return a - b.
"""
if a >= b:
return a + b
else:
return a - b
# Create a vectorized version of foo
vecfoo = np.vectorize(foo)
vecfoo(np.arange(5), 5)
array([-5, -4, -3, -2, -1])
Output type of the vectorized function
Numpy automatically evaluates the output type of the function if the otypes parameter is set to false. Here is an example to showcase it:
a = np.array([1, 2, 3, 4])
b = 2
vecfoo = np.vectorize(foo)
res = vecfoo(a, b)
print(type(res[0]))
<class 'numpy.int64'>
We can also control the output of the vectorized function by enforcing the data type of the returned value. Here is an example of how to do it.
a = np.array([1, 2, 3, 4])
b = 2
vecfoo = np.vectorize(foo, otypes=[float])
res = vecfoo(a, b)
print(type(res[0]))
<class 'numpy.float64'>
Caching in Numpy vectorization
We have already seen that if the optypes are not specified the function will call the first argument of the input to determine the number of inputs. This result can be cached thus preventing the function from running the same operation again and again. However, cache implementation slows down the subsequent calls and must only be used if the function evaluation is computationally expensive. The cache can be set by setting the parameter cache to True
Final Remarks
np.vectorize() is not the only place where we use vectorization, we have been using it all along in every day NumPy operation. Let us addition for example:
np.arange(5) + 4
array([4, 5, 6, 7, 8])
The same rule applies to different primitive functions like subtraction, multiplication, sin, cos, etc. These functions have built-in vectorization support. But our own Python version does not generally support this kind of vectorization, so we need numpy.vectorize() to do our vectorized operation in a fast and efficient manner.