Python numpy apply_along_axis用法及代码示例

本文简要介绍 python 语言中 numpy.apply_along_axis 的用法。

用法:

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

沿给定轴将函数应用于一维切片。

执行 func1d(a, *args, **kwargs) ，其中 func1d 在一维数组上运行，a 是 arr 沿轴的一维切片。

这等效于(但比)以下使用 ndindex 和 s_ ，将 ii 、 jj 和 kk 设置为索引元组：

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]

等效地，消除内部循环，这可以表示为：

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])

参数：

func1d： 函数 (M,) -> (Nj…)

这个函数应该接受一维数组。它应用于沿指定轴的一维 arr 切片。

axis： 整数

沿其切割 arr 的轴。

arr： ndarray (Ni…, M, Nk…)

输入数组。

args： 任何

func1d 的附加参数。

kwargs： 任何

func1d 的附加命名参数。

返回：

out： ndarray (Ni…, Nj…, Nk…)

输出数组。除了沿轴维度外，out 的形状与 arr 的形状相同。该轴被移除，并替换为与 func1d 的返回值形状相等的新维度。因此，如果 func1d 返回一个标量输出将比 arr 少一个维度。

例子：

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

对于返回一维数组的函数，outarr 中的维数与 arr 相同。

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

对于返回更高维数组的函数，将插入这些维度来代替轴维度。

>>> b = np.array([[1,2,3], [4,5,6], [7,8,9]])
>>> np.apply_along_axis(np.diag, -1, b)
array([[[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3]],
       [[4, 0, 0],
        [0, 5, 0],
        [0, 0, 6]],
       [[7, 0, 0],
        [0, 8, 0],
        [0, 0, 9]]])