NumPy: Best Ways to Map a Function Over an Array

In this tutorial, you’ll learn how to use NumPy to map a function over an array using different methods such as NumPy vectorize. Being able to apply the same function to each element in an array is an important skill. However, because NumPy arrays can often be quite large, we need to consider performance when mapping functions to NumPy arrays.

By the end of this tutorial, you’ll have learned:

• How to apply functions to NumPy one-dimensional and two-dimensional arrays
• How to use the NumPy vectorize function to map a function over an array
• How to use for loops and list comprehensions to apply a function to every element in a NumPy array

What is the Best Way to Map a Function to a NumPy Array

The best way to map a function to a NumPy array is to pass the array into a function directly. Not only is this the simplest way, but it is also the most readable method. The method works for arrays of any dimension.

Let’s dive into how this method works by first exploring how to map a function to a one-dimensional array in the next section.

Map a Function to a One-Dimensional NumPy Array

The simplest way to map a function to a one-dimensional array is to simply pass the array into the function. This function can be a built-in function, user defined function, or an anonymous lambda function.

Let’s take a look at creating a user-defined function that squares a number. We can then pass a NumPy array into that function to see what happens.

# Applying a User-Defined Function to NumPy Arrays
import numpy as np

def square(num):
    return num ** 2

arr = np.arange(11)
arr = square(arr)

print(arr)

# Returns: [  0   1   4   9  16  25  36  49  64  81 100]

Let’s break down what we did in the code above:

1. We imported NumPy using the alias np
2. We then defined a function, square(), which takes a single number as its input and returns the square of that number
3. We then created an array, arr, using the NumPy arange() function that holds the numbers from 0 through 10
4. We then passed that array into our function and assigned it back to our array

Similar to the example above, we can use an anonymous lambda function to map function arrays. Let’s see how we can recreate our function as a simple lambda function:

# Applying a Lambda Function to NumPy Arrays
import numpy as np

square = lambda x: x ** 2

arr = np.arange(11)
arr = square(arr)

print(arr)

# Returns: [  0   1   4   9  16  25  36  49  64  81 100]

The code in the example above behaves similar to using a user-defined function. However, we use a lambda function instead. Now that you know how to map functions to one-dimensional NumPy arrays, let’s see how we can do this for two-dimensional arrays.

Map a Function to a Two-Dimensional NumPy Array

Applying a function to each element in a two-dimensional NumPy array works in the same way as with one-dimensional arrays. We can simply pass the array into a function and map each element using that function. Let’s see how this works:

# Applying a Lambda Function to 2-D NumPy Arrays
import numpy as np

square = lambda x: x ** 2

arr = np.arange(10).reshape(2, 5)
arr = square(arr)

print(arr)

# Returns: 
# [[ 0  1  4  9 16]
#  [25 36 49 64 81]]

We can see from the code block above that the function is applied to each element in the array. This approach is quite Pythonic and allows us to be very intentional with our code.

How to Use NumPy vectorize to Map a Function to an Array

The NumPy vectorize() function is a convenience function provided by NumPy to create functions that can be applied to NumPy arrays. The function converts another function in order to apply it NumPy arrays.

It should be noted, that the function isn’t designed for performance. It simply loops over each element in the array and processes the item. Because of this, this isn’t the most recommended approach. Let’s see how we can vectorize our earlier function and apply it to each element in an array.

# Using NumPy vectorize() to Map a Function to an Array
import numpy as np

square = lambda x: x ** 2
vectorized_square = np.vectorize(square)

arr = np.arange(11)
arr = vectorized_square(arr)

print(arr)

# Returns: [  0   1   4   9  16  25  36  49  64  81 100]

In essence, this does not do much beyond what our previous examples. It actually adds a step to our processing of data and could lead to some unexpected behavior with more complex functions.

How to Use Python Map to Map a Function to a NumPy Array

Python comes with a built-in function for mapping functions to iterable items, map(). The function allows you to easily apply functions to each element in an iterable without needing to explicitly write a for loop. Let’s see how we can use the function to apply a function to each element in an array:

# Using Python map() to Map a Function to a NumPy Array
import numpy as np

arr = np.arange(11)
square = lambda x: x ** 2
mapped = np.array(list(map(square, arr)))
print(mapped)

# Returns: [  0   1   4   9  16  25  36  49  64  81 100]

In the code block above, we created an array and a lambda function. We then passed the function and array into the Python map() function, respectively. Because the map() function returns a map object, we need to convert it back into a NumPy array using the np.array() constuctor.

How to Use For Loops to Map a Function to a NumPy Array

In this section, you’ll learn how to use a Python for loop to map a function to each item in a NumPy array. This method explicitly does what the NumPy vectorize function does and may be clearer to read for readers of your code.

Let’s see how we can do this:

# Using a For Loop to Map a Function to a NumPy Array
import numpy as np

arr = np.arange(11)
square = lambda x: x ** 2

for idx in range(len(arr)):
    arr[idx] = square(arr[idx])

print(arr)

# Returns: [  0   1   4   9  16  25  36  49  64  81 100]

In the code block above, we first created our array and the function we want to apply. We then looped over each element in the array and appended it to a new array. We manipulate each element in the array by accessing the array value directly in the for loop. We assign the value to the mapped value returned by our function.

We can simplify this code by making use of a list comprehension, as you’ll learn in the following section.

How to Use List Comprehensions to Map a Function to a NumPy Array

We can simplify the method above by using a list comprehension to replace our for loop. We can use a list comprehension to create a new list of values where each item is mapped to a function. In order to turn this back into a NumPy array, we need to pass it into the np.array() constructor function.

# Using a List Comprehension to Map a Function to a NumPy Array
import numpy as np

arr = np.arange(11)
square = lambda x: x ** 2

arr = np.array([square(val) for val in arr])

print(arr)

# Returns: [  0   1   4   9  16  25  36  49  64  81 100]

Let’s break down what we did in the code block above:

1. We created our array, arr
2. We declared our lambda function to square values
3. We then used a list comprehension to iterate over each item in the array and apply the function to each item
4. We then pass the list into the np.array() function to return a NumPy array

Conclusion

In this guide, you learned how to map a function to a NumPy array. You first learned how to simply pass an array into a function, both for one-dimensional and two-dimensional arrays. Then, you learned how to use the NumPy vectorize() function to transform other functions to vector functions. From there, you learned how to use the Python map() function, for loops and list comprehension to map a function to a NumPy array.

Additional Resources

To learn more about related topics, check out the tutorials below:

• Indexing and Slicing NumPy Arrays: A Complete Guide
• How to Normalize NumPy Arrays
• NumPy where: Process Array Elements Conditionally
• NumPy vectorize: Official Documentation