如何向NumPy矩阵所有元素应用带双参数的Python函数？

问题场景 #

自定义函数定义：

def my_func(a,b):
    res = a[0] + a[1]*b
    return res

单个元素调用正常：

import numpy as np
mydata = np.matrix([[1, 2], [3, 4]])
my_par = np.array([1, 2])
res = my_func(my_par,mydata[1,1])  # 正常返回9

但使用np.vectorize批量处理矩阵所有元素时触发错误：

myfunc_vec = np.vectorize(my_func)
res = myfunc_vec(my_par,mydata)

报错信息：

in my_func
    res = a[0] + a[1]*b
IndexError: invalid index to scalar variable.

错误原因 #

np.vectorize默认会对所有输入参数做元素级拆解，它把my_par的每个元素（1和2）分别传给my_func的a参数，导致a变成了无法索引的标量，因此触发IndexError。

解决方案 #

方案1：给np.vectorize指定signature参数 #

通过signature明确定义输入参数的形状，告诉vectorize第一个参数是长度为n的一维数组，第二个参数是标量，返回值是标量，避免拆解my_par：

import numpy as np

def my_func(a,b):
    res = a[0] + a[1]*b
    return res

mydata = np.matrix([[1, 2], [3, 4]])
my_par = np.array([1, 2])

# signature格式说明：(n)表示第一个参数是一维数组，()表示第二个参数是标量，最后一个()表示返回标量
myfunc_vec = np.vectorize(my_func, signature='(n),()->()')
res = myfunc_vec(my_par, mydata)
print(res)

输出结果：

[[3 5]
 [7 9]]

方案2：直接利用NumPy广播机制（推荐） #

np.vectorize本质是Python循环包装，效率较低。NumPy原生支持广播运算，直接通过元素级运算即可实现需求，无需使用vectorize：

import numpy as np

mydata = np.matrix([[1, 2], [3, 4]])
my_par = np.array([1, 2])

res = my_par[0] + my_par[1] * mydata
print(res)

输出结果：

[[3 5]
 [7 9]]

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