Numpy中如何应用返回数组而非标量的函数？

问题背景 #

Numpy的apply_along_axis/apply_over_axes常被认为只能处理返回标量的函数，但实际场景中我们需要调用返回数组（为原数组新增维度）的函数，具体场景如下：

示例说明 #

原始简化函数：

def my_func(k):
    x = np.arange(3)
    y = x ** k
    return y

原始输入数组：

array([[1],
       [2],
       [3]])

预期将函数应用到每一行后，得到新增维度的结果：

array([[ 0,  1,  2,  3],
       [ 0,  1,  4,  9],
       [ 0,  1,  8, 27]], dtype=int32)

实际场景更新 #

实际使用的是仅支持标量输入、不支持广播的函数，例如：

def my_func(k):
    return np.random.randint(1, 4, 5) + k

可行解决方案 #

1. 直接使用np.apply_along_axis #

apply_along_axis并非只能处理返回标量的函数，当函数返回固定长度的数组时，它会自动将所有结果拼接成更高维度的数组。只需将每行的标量传入函数即可：

import numpy as np

def my_func(k):
    return np.random.randint(1, 4, 5) + k

arr = np.array([[1], [2], [3]])
# 沿axis=1遍历每行，取出标量传入函数
result = np.apply_along_axis(lambda row: my_func(row[0]), axis=1, arr=arr)
# 结果形状为(3, 5)，符合新增维度的需求

2. 用np.vectorize指定签名实现向量化 #

针对仅接收标量的函数，可以使用np.vectorize并通过signature参数明确输入输出的形状，实现对整个数组的向量化处理：

import numpy as np

def my_func(k):
    return np.random.randint(1, 4, 5) + k

arr = np.array([[1], [2], [3]])
# signature参数定义：输入为标量，输出为长度5的数组
vec_func = np.vectorize(my_func, signature='()->(5)')
result = vec_func(arr).squeeze(axis=1)  # 挤压多余维度，得到(3,5)的结果

3. 手动循环+预分配数组（极致性能） #

如果追求最高运行效率，手动循环并预先分配结果数组的方式通常比apply类方法更快，尤其适合大规模数组：

import numpy as np

def my_func(k):
    return np.random.randint(1, 4, 5) + k

arr = np.array([[1], [2], [3]])
# 预分配结果数组，指定形状和数据类型
result = np.empty((arr.shape[0], 5), dtype=np.int32)

for i in range(arr.shape[0]):
    result[i] = my_func(arr[i, 0])

内容的提问来源于stack exchange，提问作者Denis Kazakov