Numpy如何沿指定轴应用自定义函数 同时兼容1D与2D数组输入

问题描述 #

我有两个数组arr1和arr2：

arr1 = [0, 1, 2]

arr2 = [
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8],
]

我定义了如下对数组元素进行处理的自定义函数：

def func(arr):
    new_arr = arr.copy()
    new_arr[0] = new_arr[0] * 2
    new_arr[1] = new_arr[1] * 10
    new_arr[2] = new_arr[2] * 100
    return new_arr

现在我需要将该函数向量化，使其可以同时适配arr1和arr2的输入，预期输出如下：

func(arr1)
# returns [0, 10, 200]
func(arr2)
# returns
# [0, 10, 200],
# [6, 40, 500],
# [12, 70, 800],

np.vectorize无法达到预期效果，因为它会拆解数组的每一个元素单独处理，而我希望仅沿第一个轴应用该函数；np.apply_along_axis基本可用，但它无法将1D数组输入作为单个参数处理。请问最优的实现方案是什么？

解决方案 #

方案1：广播实现（性能最优，推荐） #

针对当前的计算逻辑，直接使用numpy广播特性即可实现，无需任何循环，性能最高：

import numpy as np

def func(arr):
    arr = np.asarray(arr)
    coeff = np.array([2, 10, 100])
    return arr * coeff

该实现同时兼容1D和2D输入：

• 输入1D数组时，直接对应元素相乘，输出符合预期的[0, 10, 200]
• 输入2D数组时，系数数组会自动广播到和输入相同的shape，逐行计算，输出符合预期的2D结果

方案2：兼容1D输入的apply_along_axis实现（适配复杂自定义逻辑） #

如果实际业务逻辑更复杂，无法直接改写为矩阵运算，可以通过维度判断适配np.apply_along_axis的1D输入问题：

import numpy as np

def func(arr):
    arr = np.asarray(arr)
    # 1D输入直接执行自定义逻辑
    if arr.ndim == 1:
        new_arr = arr.copy()
        new_arr[0] *= 2
        new_arr[1] *= 10
        new_arr[2] *= 100
        return new_arr
    # 2D及以上维度沿最后一个轴应用自定义逻辑
    return np.apply_along_axis(lambda x: [x[0]*2, x[1]*10, x[2]*100], axis=-1, arr=arr)

内容的提问来源于stack exchange，提问作者Zahin Zaman