NumPy向量化：如何将双参数函数应用于矩阵所有元素对

问题场景 #

现有两个NumPy数组、一个接收两个入参返回单个值的自定义函数：

import numpy as np
nparr1 = np.array([1,2])
nparr2 = np.array([[0.4, 0.8, 1.2, 1.6, 2]])

def fun(a,b):
    # 加法仅为示例，可替换为任意双参数计算逻辑
    return a + b

需要输出形状为(2,5)的结果矩阵，第i行第j列元素为fun(nparr1[i], nparr2[0,j])，即：

result = [
    [fun(1,0.4), fun(1,0.8), fun(1,1.2), fun(1,1.6), fun(1,2)],
    [fun(2,0.4), fun(2,0.8), fun(2,1.2), fun(2,1.6), fun(2,2)]
]

已验证运行np.array(np.meshgrid(nparr1, nparr2)).T可得到形状为(2,5,2)的中间配对矩阵，最后一维的两个值恰好是每一组待传入函数的参数对：

intermidiate_matrix = [[[1.  0.4]
  [1.  0.8]
  [1.  1.2]
  [1.  1.6]
  [1.  2. ]]

 [[2.  0.4]
  [2.  0.8]
  [2.  1.2]
  [2.  1.6]
  [2.  2. ]]]

实现方法 #

基于已有中间矩阵的实现 #

拿到三维中间矩阵后，使用np.apply_along_axis沿最后一个轴批量调用函数即可：

intermidiate_matrix = np.array(np.meshgrid(nparr1, nparr2)).T
result = np.apply_along_axis(lambda pair: fun(pair[0], pair[1]), axis=2, arr=intermidiate_matrix)

运行输出结果与预期完全匹配：

array([[1.4, 1.8, 2.2, 2.6, 3. ],
       [2.4, 2.8, 3.2, 3.6, 4. ]])

基于np.vectorize的简洁实现 #

之前使用np.vectorize未成功，核心原因是没有匹配数组广播规则，不需要手动生成meshgrid，调整数组维度后直接传入向量化包装的函数即可：

vectorized_fun = np.vectorize(fun)
# 将nparr1转为(2,1)列向量，与形状(1,5)的nparr2自动广播完成元素配对
result = vectorized_fun(nparr1[:, np.newaxis], nparr2)

该写法无需生成三维中间矩阵，内存占用更低，输出结果与上述方法一致。

补充说明：np.vectorize本质是Python层循环封装，不会带来计算性能提升，仅用于简化批量调用写法。如果自定义逻辑可以用NumPy原生算子实现，直接利用广播计算性能最优，比如示例中的加法可直接写为nparr1[:, np.newaxis] + nparr2，运行速度比循环类实现快1~2个数量级。

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