NumPy Functional programming: apply_along_axis() function

numpy.apply_along_axis() function

The apply_along_axis() function is used for apply a function to 1-D slices along the given axis.

Execute func1d(a, *args) where func1d operates on 1-D arrays and a is a 1-D slice of arr along axis.

This is equivalent to (but faster than) the following use of ndindex and s_, which sets each of ii, jj, and kk to a tuple of indices:

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
for ii in ndindex(Ni):
    for kk in ndindex(Nk):
        f = func1d(arr[ii + s_[:,] + kk])
        Nj = f.shape
        for jj in ndindex(Nj):
            out[ii + jj + kk] = f[jj]

Equivalently, eliminating the inner loop, this can be expressed as:

Ni, Nk = a.shape[:axis], a.shape[axis+1:]
for ii in ndindex(Ni):
    for kk in ndindex(Nk):
        out[ii + s_[...,] + kk] = func1d(arr[ii + s_[:,] + kk])

Syntax:

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

Version: 1.15.0

Parameter:

Returns: out : ndarray (Ni…, Nj…, Nk…)

The output array. The shape of out is identical to the shape of arr, except along the axis dimension.
This axis is removed, and replaced with new dimensions equal to the shape of the return value of func1d.
So if func1d returns a scalar out will have one fewer dimensions than arr.

NumPy.apply_along_axis() method Example-1:

>>> import numpy as np
>>> def my_func(x): return (x[0] + x[-1]) * 0.5 # """Average first and last element of a 1-D array"""
>>> y = np.array([[2,4,6], [1,3,5], [9,7,8]])
>>> np.apply_along_axis(my_func, 0, y)

Output:

array([5.5, 5.5, 7. ])

NumPy.apply_along_axis() method Example-2:

>>> import numpy as np
>>> def my_func(x): return (x[0] + x[-1]) * 0.5  #"""Average first and last element of a 1-D array"""
>>> y = np.array([[2,4,6], [1,3,5], [9,7,8]])
>>> np.apply_along_axis(my_func, 1, y)

Output:

array([4. , 3. , 8.5])

For a function that returns a 1D array, the number of dimensions in outarr is the same as arr.

NumPy.apply_along_axis() method Example-3:

>>> import numpy as np
>>> y = np.array([[2,4,6], [1,3,5], [9,7,8]])
>>> np.apply_along_axis(sorted, 1, y)

Output:

array([[2, 4, 6],
       [1, 3, 5],
       [7, 8, 9]])

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

NumPy.apply_along_axis() method Example-4:

>>> import numpy as np
>>> y = np.array([[2,4,6], [1,3,5], [9,7,8]])
>>> np.apply_along_axis(np.diag, -1, y)

Output:

array([[[2, 0, 0],
        [0, 4, 0],
        [0, 0, 6]],

       [[1, 0, 0],
        [0, 3, 0],
        [0, 0, 5]],

       [[9, 0, 0],
        [0, 7, 0],
        [0, 0, 8]]])

Python - NumPy Code Editor:

Previous: NumPy Home
Next: apply_over_axes() function