NumPy apply_along_axis()

import numpy as np

# create a 2D array
arr = np.array([[1, 2, 3],
                [4, 5, 6],
                [7, 8, 9]])

# function to calculate the sum of an array
def sumArray(arr):
    return np.sum(arr)

# apply the sumArray function along the rows (axis=1)
result = np.apply_along_axis(sumArray, axis=1, arr=arr)

print(result)

# Output: [ 6 15 24]

apply_along_axis() Syntax

The syntax of apply_along_axis() is:

numpy.apply_along_axis(func1d, axis, arr, *args, **kwargs)

apply_along_axis() Arguments

The apply_along_axis() method takes following arguments:

• func1d - the function to apply along the specified axis
• axis - the axis along which the function is applied
• arr - the input array to which the function will be applied
• *args and **kwargs - additional arguments and keyword arguments present in func1d

apply_along_axis() Return Value

The apply_along_axis() method returns the resultant array with functions applied.

Example 1: Apply a Function That Returns a Single Value

import numpy as np

# create a 2D array
arr = np.array([[1, 2, 3],
                [4, 5, 6],
                [7, 8, 9]])

# define a function to return the last element of an array
def lastItem(subArr):
    return np.max(subArr[-1])

# return last item along the rows (axis=1)
result = np.apply_along_axis(lastItem, axis=1, arr=arr)
print(result)

# return last item along the columns (axis=0)
result = np.apply_along_axis(lastItem, axis=0, arr=arr)
print(result)

Output

[3 6 9]
[7 8 9]

Example 2: Apply a Function That Returns an Array of Values

We can also return an array of values from the function.

import numpy as np

# create a 2D array
arr = np.array([[1, 2, 3],
               [4, 5, 6],
              [7, 8, 9]])

# function to return the square of elements of an array
def square(arr):
    return (arr*arr)

# return the square of elements
result = np.apply_along_axis(square, axis = 0, arr=arr
)
print(result)

Output

[[ 1  4  9]
 [16 25 36]
 [49 64 81]]

Example 3: Apply a Function That Returns an N-D Array of Values

We can return an n-D array of values from the function.

Let's looks at an example.

import numpy as np

# create a 2D array
arr = np.array([[1, 2, 3],
               	[4, 5, 6],
                [7, 8, 9]])

# define a function that returns a 2D array
def square_and_cube(arr):
    return np.array([arr**2, arr**3])

# apply the square_and_cube function along the columns (axis=0)
result = np.apply_along_axis(square_and_cube, axis=0, arr=arr)

print('Along axis 0\n',result)

# apply the square_and_cube function along the rows (axis=1)
result = np.apply_along_axis(square_and_cube, axis=1, arr=arr)

print('Along axis 1\n',result)

Output

Along axis 0
[[[  1   4   9]
  [ 16  25  36]
  [ 49  64  81]]

 [[  1   8  27]
  [ 64 125 216]
  [343 512 729]]]

Along axis 1
 [[[  1   4   9]
  [  1   8  27]]

 [[ 16  25  36]
  [ 64 125 216]]

 [[ 49  64  81]
  [343 512 729]]]

For a function that returns a higher dimensional array, the dimensions are inserted in place of the axis dimension.

Example 4: Apply a lambda Function to an Array

Instead of defining a function, we can directly pass a lambda function.

Let's look at an example.

import numpy as np

# create a 2D array
arr = np.array([[1, 2, 3],
               	[4, 5, 6],
              	[7, 8, 9]])

# apply the summation lambda function along the rows (axis=1)
result = np.apply_along_axis(lambda arr:np.sum(arr), axis=1, arr=arr)

print(result)

Output

[ 6 15 24]

Recommended Reading

Numpy apply_over_axes()