如何用np.vectorize对3D numpy数组按指定签名向量化函数？

你需要将函数f(x, y)应用到形状为(N,M,2)的numpy数组a上，数组最后一个维度的两个元素分别作为x和y传入函数。以下是示例场景：

示例数组与函数 #

import numpy as np

a = np.array([[[1, 1],
        [2, 1],
        [3, 1]],

       [[1, 2],
        [2, 2],
        [3, 2]],

       [[1, 3],
        [2, 3],
        [3, 3]]])

def function_to_vectorize(x, y):
   # 示例函数，无实际业务意义
   if x>2 and y-x>0:
      sum = 0
      for i in range(y):
         sum+=i
      return sum
   else:
      sum = y
      for i in range(x):
         sum-=i
      return sum

期望的应用效果（嵌套列表推导式实现） #

result = [[function_to_vectorize(element[0], element[1]) for element in vector] for vector in a]
# 输出结果：
# array([[ 1,  0, -2],
#        [ 2,  1, -1],
#        [ 3,  2,  3]])

用np.vectorize实现的方法 #

步骤说明 #

1. 包装向量化函数：使用np.vectorize将原标量函数转换为可处理数组输入的版本。
2. 拆分输入数组：将原数组a按最后一个维度拆分，得到分别对应x和y的两个(N,M)形状数组。
3. 调用向量化函数：将拆分后的数组传入向量化函数，直接得到目标结果。

完整代码示例 #

import numpy as np

# 定义原函数
def function_to_vectorize(x, y):
   if x>2 and y-x>0:
      sum = 0
      for i in range(y):
         sum+=i
      return sum
   else:
      sum = y
      for i in range(x):
         sum-=i
      return sum

# 创建向量化函数，可指定输出类型避免推断错误
vectorized_func = np.vectorize(function_to_vectorize, otypes=[int])

# 拆分数组的x、y分量
x = a[..., 0]  # 取最后一个维度的第0个元素，形状(N,M)
y = a[..., 1]  # 取最后一个维度的第1个元素，形状(N,M)

# 调用向量化函数得到结果
result = vectorized_func(x, y)
print(result)
# 输出：
# [[ 1  0 -2]
#  [ 2  1 -1]
#  [ 3  2  3]]

注意事项 #

• np.vectorize本质是包装了循环逻辑，并非底层真正的向量化实现，性能与列表推导式接近，适合快速实现逻辑而非极致性能优化。
• 如果函数返回固定类型，通过otypes参数指定输出类型能避免潜在的类型推断错误。

内容的提问来源于stack exchange，提问作者Salvatore Daniele Bianco